"""
benchmarks/common/episodic.py -- the closed-loop counterpart to
`benchmarks/common/harness.py::run_harness()`.

`run_harness()` is a **static-dataset** loop: filter a frozen item list, call
a backend once per item, score against a gold label, write a bundle. ViZDoom
(the first, and so far only, consumer of this module -- PRD.md 8.1d) is
**episodic**: the item sequence is generated by the environment *in response
to the model's own actions*, and there is no gold label to score against at
all. Threading an `if episodic:` branch through a function that
`benchmarks/jevbench`, `benchmarks/jabr_v2` and `benchmarks/screenspot_v2`
all depend on staying exactly as it is would be the wrong shape -- so this is
a new, separate function instead (PRD.md 8.1d, "Harness architecture: a new
episodic loop, not a bent static one").

Reused, unchanged, from the static-dataset path: `build_backend()` /
`make_arg_parser()` (`benchmarks/common/harness.py`), `write_bundle()`
(`benchmarks/common/evidence.py`), `eval/metrics.py::full_report()` (used
here against the rubric's own prescribed action, not a dataset gold label --
see `rubric_agreement` below and PRD.md 8.1d Finding 4).

One call to `run_episodic_harness()` covers ONE (scenario, rubric) pair
across every one of that scenario's seeds -- `benchmarks/vizdoom/run.py`'s
CLI loops over scenarios/rubrics and calls this once per combination, the
same way Von's own `run_doom_benchmark.py` loops over its two scenarios
without folding them into a single call.
"""

from __future__ import annotations

import json
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Callable, Optional

import numpy as np

from benchmarks.common.evidence import write_bundle
from eval.metrics import full_report


# --- Policy: the one seam between "text/options in, choice out" backends ---
# and an episodic loop that also has a raw environment snapshot available.

class BackendPolicy:
    """Adapts any `benchmarks.common.backends` `ChoiceBackend` -- which only
    ever sees an item's rendered text and options, never a raw environment
    snapshot -- into the `Policy` interface `run_episodic_harness()` actually
    drives. This is the *only* place the widening from "text in, choice out"
    to "snapshot, text, choice out" happens, specifically so a text-only
    `ChoiceBackend` never has to know episodic benchmarks exist."""

    def __init__(self, backend: Any):
        self.backend = backend

    def describe(self) -> dict:
        return self.backend.describe()

    def act(self, *, snapshot: Any, state_text: str, options: list[str], instructions: Optional[str]) -> dict:
        return self.backend.predict_choice(state_text, options, instructions=instructions)


class OraclePolicy:
    """PRD.md 8.1d Finding 4: the rubric is a deterministic function of the
    environment snapshot, so it can be run directly as a zero-cost, no-model,
    no-GPU baseline policy. Its "prediction" always equals its own rule by
    construction (its `rubric_agreement` is trivially 1.0) -- this policy
    exists to measure real episode outcomes (kills / survival) for the
    oracle baseline (PRD.md 8.1d checklist step 3: "the first real number,
    needs no model, no GPU, no checkpoint"), not to be meaningfully scored
    for rubric agreement with itself. `run_episodic_harness()` still computes
    and reports that number for output-shape consistency with every other
    policy, and its being exactly 1.0 is itself a useful sanity check that
    the scoring path works at all."""

    name = "rubric_oracle"

    def __init__(self, rubric_action_fn: Callable[[Any, str], str], scenario_name: str):
        self._fn = rubric_action_fn
        self._scenario = scenario_name

    def describe(self) -> dict:
        return {
            "backend": self.name,
            "note": "NOT a model -- runs the rubric's own deterministic rule directly as a "
                    "policy. See PRD.md 8.1d Finding 4 and benchmarks/vizdoom/rubric.py.",
        }

    def act(self, *, snapshot: Any, state_text: str, options: list[str], instructions: Optional[str]) -> dict:
        action = self._fn(snapshot, self._scenario)
        probabilities = {opt: (1.0 if opt == action else 0.0) for opt in options}
        return {"choice": action, "probabilities": probabilities, "confidence": 1.0, "latency_ms": 0.0}


# --- DAgger-teacher data emission (PRD.md 8.1d Finding 4 point 3 / 5.3b) ----

class DaggerSink:
    """Emits on-policy `(text, frame, action)` triples from every text-arm
    rollout -- the cheapest available route to the future vision arm's
    training data (PRD.md 8.1d Finding 4 point 3, wiring into 5.3b's
    already-designed DAgger data-collection method rather than reinventing
    it). Per PRD.md Section 9's boundary rule, this data is a function of
    *this checkpoint's own rollouts*, not a static licensed corpus, so it
    stays in this repo under `training/dagger_data/` -- never in the sibling
    better-jev-bench repo.

    **Storage format, this build's own call, recorded here rather than
    picked silently:**

    ```
    training/dagger_data/<benchmark>/<run_id>/
      triples.jsonl   one line per tick: scenario, seed, rubric, episode
                       index, tick, the exact `state` text passed to the
                       policy, the saved frame's path (relative to this
                       directory), the rubric's own prescribed action (the
                       "teacher"), and the policy's actual (post-fallback)
                       action.
      frames/*.png     one PNG per tick, named
                       <scenario>_<rubric>_<seed>_<episode>_<tick>.png
    ```

    JSONL, not one-JSON-object-per-file, because a future consumer will want
    to stream/shuffle rows for training -- the same shape
    `training/*.py`'s existing data pipelines already expect. PNG frames as
    separate files, not embedded/base64, because they are the expensive part
    and a training dataloader wants to `Image.open()` them directly, matching
    `training/decoder_lora_lib.py`'s own `PromptDataset` convention for image
    rows.

    **Gitignored** (`training/dagger_data/`, added to `.gitignore` alongside
    this file) -- the same discipline as `checkpoints/` and `results/*.raw`:
    real, regenerable output, not something to check into git. Nothing
    downstream consumes this data yet (there is no vision-arm trainer wired
    up to read it) -- this class exists so that the day one is built, the
    data-collection half is already done and already validated end to end,
    not so anything reads it today.
    """

    def __init__(self, out_dir: Path):
        self.out_dir = out_dir
        self.frames_dir = out_dir / "frames"
        self.frames_dir.mkdir(parents=True, exist_ok=True)
        self.jsonl_path = out_dir / "triples.jsonl"
        self._n_emitted = 0
        self._f = self.jsonl_path.open("a")

    def emit(
        self, *, scenario: str, seed: int, rubric: str, episode_index: int, tick: int,
        state_text: str, frame: Optional[np.ndarray], teacher_action: str, model_action: Optional[str],
    ) -> None:
        frame_rel = None
        if frame is not None:
            from PIL import Image  # lazy: this class is only ever instantiated when --dagger-out is set

            fname = f"{scenario}_{rubric}_{seed}_{episode_index}_{tick:04d}.png"
            Image.fromarray(frame).save(self.frames_dir / fname)
            frame_rel = f"frames/{fname}"
        row = {
            "scenario": scenario, "seed": seed, "rubric": rubric, "episode_index": episode_index,
            "tick": tick, "state": state_text, "frame_path": frame_rel,
            "teacher_action": teacher_action, "model_action": model_action,
        }
        self._f.write(json.dumps(row, sort_keys=True, default=str) + "\n")
        self._n_emitted += 1

    def close(self) -> dict:
        self._f.close()
        return {"enabled": True, "out_dir": str(self.out_dir), "jsonl": str(self.jsonl_path), "n_triples": self._n_emitted}


# --- The episodic scenario contract ----------------------------------------

@dataclass
class ScenarioSpec:
    """Everything `run_episodic_harness()` needs to run one scenario, and
    nothing benchmark-specific baked into the loop itself -- ViZDoom is the
    only implementation today (`benchmarks/vizdoom/run.py` builds these), but
    nothing here assumes ViZDoom by name."""

    name: str
    seeds: list[int]
    metric_name: str  # e.g. "kills" or "survival_s" -- becomes a results_summary key
    env_factory: Callable[[str, int, bool], Any]  # (scenario_name, seed, capture_frames) -> env
    observe_fn: Callable[[Any, str, Optional[str]], str]  # (snapshot, scenario_name, last_action) -> state text
    rubric_action_fn: Callable[[Any, str], str]  # (snapshot, scenario_name) -> the oracle's prescribed action
    extract_metric_fn: Callable[[Any, int, int], float]  # (last_snapshot, steps, tics) -> metric value


def _percentile(values: list[float], p: float) -> Optional[float]:
    if not values:
        return None
    return float(np.percentile(np.array(values, dtype=float), p))


def run_episodic_harness(
    *,
    benchmark: str,
    args: Any,
    scenario: ScenarioSpec,
    rubric: str,
    get_choice_fn: Callable[[str, str], tuple[list[str], str]],
    policy: Any,
    tics: int,
    step_cap: int,
    episodes_per_seed: int = 1,
    seeds_override: Optional[list[int]] = None,
    dataset_provenance: dict,
    dagger_sink: Optional[DaggerSink] = None,
) -> dict:
    """Runs `scenario` under `rubric` framing, driven by `policy`
    (`BackendPolicy` wrapping a real/mock/random `ChoiceBackend`, or
    `OraclePolicy`), for every seed in `scenario.seeds` (or `seeds_override`)
    times `episodes_per_seed` episodes each, and writes a PRD.md 8.2 evidence
    bundle exactly like `run_harness()` does for the static-dataset path.

    **8.1a's no-schema-repair rule, applied to a closed loop.** A static
    benchmark can simply mark an unparseable response wrong; an episodic one
    cannot -- the environment must advance by *some* action every tick
    regardless of what the policy returned. An out-of-schema response (not
    one of `options`, or an exception from `policy.act()`) is therefore
    substituted with a fixed fallback (`options[0]`) purely so the episode
    can continue, but this substitution is never silent: it is counted
    per-episode and in aggregate (`n_out_of_schema`), recorded per-tick in
    the raw evidence file (`out_of_schema: true`, the real `model_action`
    preserved alongside the `effective_action` actually sent to the
    environment), and can never itself count as agreeing with the rubric
    (the raw model_action, not the substituted one, is what's scored).

    **`rubric_agreement` is not accuracy.** ViZDoom has no gold label for the
    "right" action at any given tick -- `eval/metrics.py`'s accuracy/Brier/
    ECE are computed here against the rubric's own deterministic prescribed
    action (`scenario.rubric_action_fn`) instead, which is a well-defined but
    different question ("how often, and how confidently, does the policy do
    what the prompt told it to"). See PRD.md 8.1d Finding 4 and the note
    embedded in the returned `results_summary["rubric_agreement"]` itself --
    this caveat travels with the number, not just with this docstring.
    """
    seeds = seeds_override if seeds_override is not None else list(scenario.seeds)
    options, instructions = get_choice_fn(scenario.name, rubric)
    n_options = len(options)

    probs_rows: list[list[float]] = []
    label_idx_rows: list[int] = []
    raw_rows: list[dict] = []
    per_episode: list[dict] = []
    n_out_of_schema = 0

    for seed in seeds:
        env = scenario.env_factory(scenario.name, seed, dagger_sink is not None)
        for ep_idx in range(episodes_per_seed):
            snap = env.reset()
            last_action: Optional[str] = None
            last_snap = snap
            steps = 0
            latencies: list[float] = []
            ep_out_of_schema = 0

            while not env.is_finished() and steps < step_cap:
                state_text = scenario.observe_fn(snap, scenario.name, last_action)
                teacher_action = scenario.rubric_action_fn(snap, scenario.name)

                error = None
                try:
                    out = policy.act(snapshot=snap, state_text=state_text, options=options, instructions=instructions)
                    model_action = out.get("choice")
                    out_of_schema = model_action not in options
                except Exception as exc:  # defensive: a future free-generation policy could raise mid-parse
                    out = {"probabilities": {}, "confidence": None, "latency_ms": None}
                    model_action = None
                    out_of_schema = True
                    error = str(exc)

                if out_of_schema:
                    n_out_of_schema += 1
                    ep_out_of_schema += 1
                    effective_action = options[0]  # environment must still advance -- never scored as agreement
                else:
                    effective_action = model_action

                if out.get("latency_ms") is not None:
                    latencies.append(out["latency_ms"])

                probabilities = out.get("probabilities") or {}
                probs_rows.append([probabilities.get(opt, 0.0) for opt in options])
                label_idx_rows.append(options.index(teacher_action))

                raw_rows.append({
                    "scenario": scenario.name, "rubric": rubric, "seed": seed, "episode_index": ep_idx,
                    "tick": steps, "state": state_text, "options": options,
                    "teacher_action": teacher_action, "model_action": model_action,
                    "effective_action": effective_action, "out_of_schema": out_of_schema,
                    "probabilities": probabilities, "confidence": out.get("confidence"),
                    "latency_ms": out.get("latency_ms"), "error": error,
                })

                if dagger_sink is not None:
                    dagger_sink.emit(
                        scenario=scenario.name, seed=seed, rubric=rubric, episode_index=ep_idx, tick=steps,
                        state_text=state_text, frame=getattr(snap, "frame", None),
                        teacher_action=teacher_action, model_action=effective_action,
                    )

                last_action = effective_action
                env.step(effective_action, tics=tics)
                new_snap = env.get_snapshot()
                if new_snap is not None:
                    last_snap, snap = new_snap, new_snap
                steps += 1

            metric_value = scenario.extract_metric_fn(last_snap, steps, tics)
            per_episode.append({
                "seed": seed, "episode_index": ep_idx, "steps": steps,
                scenario.metric_name: metric_value,
                "n_out_of_schema": ep_out_of_schema,
                "latency_ms_p50": _percentile(latencies, 50),
                "latency_ms_p95": _percentile(latencies, 95),
            })
        env.close()

    if probs_rows:
        report = full_report(np.array(probs_rows), np.array(label_idx_rows), n_options)
    else:
        report = {"n": 0, "accuracy": None, "brier": None, "ece": None}

    values = [ep[scenario.metric_name] for ep in per_episode]
    metric_aggregate = {
        "mean": float(np.mean(values)) if values else None,
        "sd": float(np.std(values)) if values else None,
        "min": float(np.min(values)) if values else None,
        "max": float(np.max(values)) if values else None,
        "n_episodes": len(values),
    }

    policy_info = policy.describe()
    results_summary = {
        "scenario": scenario.name,
        "rubric": rubric,
        "tics": tics,
        "step_cap": step_cap,
        "episodes_per_seed": episodes_per_seed,
        "seeds": seeds,
        "n_decisions": len(raw_rows),
        "n_out_of_schema": n_out_of_schema,
        f"{scenario.metric_name}_by_episode": per_episode,
        f"{scenario.metric_name}_aggregate": metric_aggregate,
        "rubric_agreement": {
            **report,
            "note": (
                "NOT accuracy against a gold label -- ViZDoom has no gold label for the 'right' "
                "action at any tick. This measures how often, and how confidently, the policy's "
                "choice matches the rubric's own deterministic prescribed action "
                "(benchmarks/vizdoom/rubric.py::rubric_action). See PRD.md 8.1d Finding 4."
            ),
        },
        "dagger": dagger_sink.close() if dagger_sink is not None else {"enabled": False},
    }

    schema_summary = {"options": options, "n_options": n_options, "instructions": instructions}
    model_info = {"instantiated": True, **policy_info}

    policy_slug = str(policy_info.get("backend", "policy"))
    run_id = (
        f"{benchmark}-{scenario.name}-{rubric}-{policy_slug}-"
        f"{time.strftime('%Y%m%dT%H%M%SZ', time.gmtime())}"
    )
    manifest_path, raw_path = write_bundle(
        out_dir=Path(args.out_dir),
        run_id=run_id,
        benchmark=benchmark,
        seed=args.seed,
        dataset_provenance=dataset_provenance,
        schema_filter_summary=schema_summary,
        model_info=model_info,
        results_summary=results_summary,
        raw_rows=raw_rows,
        extra={"episodic": True},
    )
    print(f"wrote {manifest_path}")
    print(f"wrote {raw_path}")
    print(json.dumps(results_summary, indent=2, default=str))
    return results_summary