"""Local evaluation: run the hybrid agent on public games and report
per-game scores in the spirit of the official RHAE metric.

    per-level = min(human/agent, 1.0) ^ 2      (human baseline unknown locally;
                                                we report completion + actions)
    game      = level-weighted average (1-indexed weights)
    total     = average over games

Without official human baselines this is a RELATIVE tool: compare agent
variants against each other, never treat absolute numbers as leaderboard
predictions.
"""

from __future__ import annotations

import argparse
import json
import sys
import time
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

import yaml  # noqa: E402

from data.collect_selfplay import PUBLIC_GAMES_FALLBACK, grids_from, make_env, meta_from  # noqa: E402
from data.mock_env import step_env  # noqa: E402
from qwenjev.hybrid import HybridAgent, HybridConfig  # noqa: E402
from qwenjev.system1 import QwenJevModel, System1Config  # noqa: E402
from qwenjev.system2 import System2Config, System2Controller  # noqa: E402
from qwenjev.vision import TinyVLMConfig  # noqa: E402


def run_game(env, game_id: str, brain: HybridAgent, max_actions: int) -> dict:
    """Play one episode against any environment (mock or official).

    Lifecycle: WIN ends the episode; GAME_OVER triggers the hybrid brain's
    episode-boundary hook and a fresh reset, then play continues within the
    same action budget — identical budget for every agent variant, so
    best_level/actions stay comparable across runs."""
    env.reset()
    actions = 0
    resets = 0
    best_level = 0
    final_state = "NOT_FINISHED"
    started = time.time()
    for _ in range(max_actions):
        meta = meta_from(env)
        final_state = meta["state"]
        if final_state == "WIN":
            break
        if final_state == "GAME_OVER":
            brain.on_game_over(game_id, [f"game over at action {actions}"])
            env.reset()
            resets += 1
            continue
        grids = grids_from(env)
        best_level = max(best_level, meta["levels_completed"])
        cand = brain.decide(
            game_id, grids,
            levels_completed=meta["levels_completed"],
            win_levels=meta["win_levels"],
            available_actions=meta["available_actions"],
            game_state=meta["state"],
        )
        if cand.is_complex:
            step_env(env, cand.name, int(cand.x or 0), int(cand.y or 0))
        else:
            step_env(env, cand.name)
        actions += 1
        best_level = max(best_level, meta_from(env)["levels_completed"])
    return {
        "game_id": game_id,
        "win": final_state == "WIN",
        "final_state": final_state,
        "best_level": best_level,
        "resets": resets,
        "actions": actions,
        "seconds": round(time.time() - started, 1),
    }


def build_s1(cfg: dict):
    """System-1 from the yaml config; checkpoint loaded when present."""
    s1_cfg = cfg.get("system1", {})
    from pathlib import Path

    # v0.4.9: pool_vision auto-detection. The saved qwenjev_s1.pt config
    # carries the flag from v0.4.9 on; ABSENCE = v3-era checkpoint =
    # last-position pooling (heads built with matching H input).
    pool_vision = s1_cfg.get("pool_vision")
    if pool_vision is None:
        pool_vision = False
        try:
            import torch as _torch

            ck = (Path(cfg.get("out_dir", "checkpoints/qwenjev-s1"))
                  / "qwenjev_s1.pt")
            if ck.exists():
                payload = _torch.load(ck, map_location="cpu",
                                      weights_only=False)
                pool_vision = bool(
                    (payload.get("config") or {}).get("pool_vision", False))
        except Exception:
            pool_vision = False
    model_cfg = System1Config(
        backend=s1_cfg.get("backend", "tiny"),
        hf_model=s1_cfg.get("hf_model", "Qwen/Qwen3-VL-2B-Instruct"),
        max_leaf_tokens=int(s1_cfg.get("max_leaf_tokens", 512)),
        score_levels=int(s1_cfg.get("score_levels", 10)),
        leaf_batch_size=int(s1_cfg.get(
            "eval_leaf_batch", s1_cfg.get("leaf_batch_size", 32))),
        pool_vision=bool(pool_vision),
        tiny=TinyVLMConfig(**(s1_cfg.get("tiny") or {})),
    )
    s1 = QwenJevModel(model_cfg).build()
    ckpt = Path(cfg.get("out_dir", "checkpoints/qwenjev-s1")) / "qwenjev_s1.pt"
    if ckpt.exists():
        s1.load_heads(ckpt.parent)
        # stage B/C continuity: merge the SFT adapters into the fp16 base
        # so evaluation runs on the SAME hidden states the heads saw
        if model_cfg.backend == "hf":
            from qwenjev.hf_compat import resume_adapters

            hf_lora = ckpt.parent / "hf_lora"
            if hf_lora.exists():
                s1._backbone, _merged = resume_adapters(
                    s1._backbone, str(hf_lora), merge=True)
                print(f"stage B LoRA merged -> {hf_lora}")
        print(f"System-1 checkpoint loaded -> {ckpt}")
    return s1


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--config", default="configs/train.yaml")
    ap.add_argument("--games", nargs="*", default=None)
    ap.add_argument("--out", default="eval_report.json")
    args = ap.parse_args()
    cfg = yaml.safe_load(open(args.config))

    s1 = None
    try:
        s1 = build_s1(cfg)
    except Exception as exc:
        print(f"[warn] System-1 unavailable ({exc}); heuristic policy will run")

    s2 = None
    if cfg.get("system2", {}).get("enabled"):
        s2 = System2Controller(System2Config(endpoint=cfg["system2"].get("endpoint")))

    brain = HybridAgent(s1=s1, s2=s2, config=HybridConfig(game_time_limit_s=600))
    games = args.games or PUBLIC_GAMES_FALLBACK[:5]  # default: quick smoke set

    report = []
    max_actions = int(cfg.get("eval_max_actions", 1500))
    for game_id in games:
        try:
            env = make_env(game_id, None)
            res = run_game(env, game_id, brain, max_actions=max_actions)
            report.append(res)
            print(res)
        except Exception as exc:
            print(f"[warn] {game_id}: {exc}")

    Path(args.out).write_text(json.dumps(report, indent=2))
    wins = sum(1 for r in report if r["win"])
    print(f"\nsummary: {wins}/{len(report)} wins -> {args.out}")


if __name__ == "__main__":
    main()