"""The episode loop everything shares, and the optional verifiers taskset.

`play_episode` is the whole game: capture, the probe, the tests and the
fixture verifier all drive it, so the thing evaluated and the thing tested
cannot drift apart. The verifiers Env is a thin adapter over the same engine.
"""

from __future__ import annotations

from typing import Any, Callable

from .engine import Engine
from .generator import MAX_TURNS, world_for_seed
from .protocol import render_cost, render_lookup, render_rejection, render_screen, system_prompt
from .reward import ROLES, WEIGHTS, metrics, score

TASKSET_ID = "alert-triage"

Responder = Callable[[str, list[dict[str, Any]], dict[str, Any]], str | None]


def opening_message(engine: Engine) -> str:
    return render_screen(engine.screen()) + "\n\n" + render_cost(engine.minutes, MAX_TURNS - engine.turns)


def observation_for(engine: Engine, step) -> str:
    """The text the model sees after a reply."""
    left = MAX_TURNS - engine.turns
    if not step.accepted:
        return render_rejection(step.reason or "rejected", left)
    if step.action["action"] == "lookup":
        return render_lookup(step.action["what"], step.result, engine.world["currency"]) + "\n\n" + render_cost(engine.minutes, left)
    return "Disposition recorded. The episode is over."


def play_episode(seed: int, respond: Responder, world: dict[str, Any] | None = None) -> dict[str, Any]:
    """Drive one alert against a callable `(prompt, transcript, view) -> reply text`.

    `view` is the structured observation (what scripted policies read);
    `transcript` is the text history (what a model reads). A `None` reply is
    recorded as a rejected turn — a dropped model call is a thing that
    happened, not a turn to discard.
    """
    world = world or world_for_seed(seed)
    engine = Engine(world)
    prompt = system_prompt()
    transcript: list[dict[str, Any]] = [{"reply": None, "action": None, "observation": opening_message(engine)}]

    while not engine.done:
        reply = respond(prompt, transcript, engine.view())
        step = engine.step(reply)
        transcript.append({"reply": reply, "action": step.action, "rejection": step.reason, "observation": observation_for(engine, step)})

    episode = engine.episode()
    return {
        "seed": seed,
        "outcome": engine.outcome,
        "truncated": episode.truncated,
        "transcript": transcript,
        "rewards": score(episode),
        "weights": WEIGHTS,
        "roles": ROLES,
        "metrics": metrics(episode),
        "info": {
            "label": world["label"],
            "template": world["template"],
            "overlay": world["overlay"],
            "typology": world["typology"],
            "planted": world["planted"],
            "reference_minutes": world["reference_minutes"],
            "reference_policy": world["reference_policy"],
            "held_out": world["held_out"],
        },
    }


def replay(seed: int, replies: list[str | None], world: dict[str, Any] | None = None) -> dict[str, Any]:
    """Re-score a recorded run from its reply strings alone. What `verify` does."""
    replies = list(replies)
    engine = Engine(world or world_for_seed(seed))
    for reply in replies:
        if engine.done:
            break
        engine.step(reply)
    episode = engine.episode()
    return {"rewards": score(episode), "metrics": metrics(episode), "outcome": engine.outcome, "truncated": episode.truncated}


def build_taskset():  # pragma: no cover - requires the optional extra
    """The verifiers v1 Taskset. Imported lazily so the module stays optional."""
    import verifiers as vf

    class AlertTriageEnv(vf.Env):
        """One agent seat, `analyst`. The engine plays the user role: every
        observation is computed, never generated, which is why the reward is
        verifiable at all."""

        async def run(self, task, agent, trace):
            world = world_for_seed(task.data.info["seed"])
            engine = Engine(world)
            async with agent.interaction(system=system_prompt()) as session:
                message = opening_message(engine)
                while not engine.done:
                    reply = await session.turn(message)
                    step = engine.step(reply)
                    message = observation_for(engine, step)
                episode = engine.episode()
                for key, value in score(episode).items():
                    trace.record_reward(key, 0.0 if value is None else value, WEIGHTS[key])
                for key, value in metrics(episode).items():
                    if value is not None:
                        trace.record_reward(key, value, 0.0)
                trace.info["label"] = world["label"]
                trace.info["outcome"] = engine.outcome

    return AlertTriageEnv
