import benchflow as bf
from benchflow import FunctionUser, RoundResult
from benchflow.rollout import RolloutConfig, Scene
from benchflow._utils.benchmark_repos import resolve_source


def progressive(round: int, instruction: str, rr: RoundResult | None) -> str | None:
    if round == 0:
        return instruction.split("\n")[0]                # terse: first line only
    if rr and (rr.rewards or {}).get("reward", 0) >= 1.0:
        return None                                      # passed, stop
    if round >= 3:
        return None                                      # cap at 3 rounds
    return (
        f"Tests failed:\n{rr.verifier_output}\n\n"       # show failures + spec
        f"Full spec:\n{instruction}"
    )


config = RolloutConfig(
    task_path=resolve_source("benchflow-ai/swebenchpro", path="instance_flipt-io__flipt-..."),
    scenes=[Scene.single(agent="opencode", model="anthropic/claude-sonnet-4-6")],
    user=FunctionUser(progressive),
    max_user_rounds=3,
    environment="daytona",
)
result = await bf.run(config)

round 0   "Fix the bug described here: <one-line summary>"
            agent attempts → tests fail
round 1   "Tests <names> failed. Here is the full requirements section: <half of spec>."
            agent retries → tests still fail
round 2   "Still failing. Here's the full original spec: <complete instruction>"
            agent makes final attempt

setup() → start() → install_agent()
    ↓
[oracle setup if oracle_access=True: read /solution, hide it from agent]
    ↓
user.setup(instruction, solution)        ← once
    ↓
┌─ user.run(round, instruction, rr) → str | None
│      │ None: break
│      ↓
│   connect_as(role)
│   execute(prompts=[prompt])
│   disconnect()
│      ↓
│   soft_verify()                         ← partial hardening, sandbox stays alive
│      ↓
│   build RoundResult, log, repeat
└─    │
      ↓ (loop ends when user returns None or max_user_rounds reached)
[oracle restore: mv /solution_oracle_backup → /solution for final verify]
    ↓
verify()                                  ← full hardening, final reward
    ↓
cleanup()

from benchflow import BaseUser, RoundResult


class MyUser(BaseUser):
    async def setup(self, instruction: str, solution: str | None = None) -> None:
        """Called once before round 0.

        instruction — the original task instruction (from instruction.md)
        solution    — gold answer if oracle_access=True, else None
        """
        self.spec = instruction
        self.gold = solution

    async def run(
        self,
        round: int,
        instruction: str,
        round_result: RoundResult | None = None,
    ) -> str | None:
        """Return the next prompt, or None to stop.

        round — 0-indexed
        instruction — original task instruction (unchanged each round)
        round_result — None on round 0; previous round's outcome on subsequent rounds
        """
        ...

@dataclass
class RoundResult:
    round: int                     # 0-indexed
    trajectory: list[dict]         # ACP events from this round only
    rewards: dict | None           # verifier rewards (None if verifier crashed)
    verifier_output: str | None    # raw verifier stdout/log
    verifier_error: str | None     # exception message if verifier failed
    n_tool_calls: int              # tool calls in this round

def fn(round, instruction, rr): ...
user = FunctionUser(fn)

async def afn(round, instruction, rr): ...
user = FunctionUser(afn)

user: BaseUser | None = None     # the callback
max_user_rounds: int = 5         # cap on rounds (loop also stops when user returns None)
oracle_access: bool = False      # expose gold solution to user.setup()

[verifier.hardening]
cleanup_conftests = false

[User] round 2: prompt='Try again, focusing on...'
ERROR  user.run() failed at round 2: KeyError: 'spec_section'