"""Per-checkpoint eval worker: watch the adapter repo, eval each new checkpoint.

Loop: poll ADAPTER_REPO for checkpoint-N dirs without eval results, hot-load each
as a LoRA into the already-running vLLM (runtime LoRA updates — no server restart),
run the held-out eval on a fixed task subset, and upload per-checkpoint results
plus a rolling eval/summary.json (the live learning curve).

Base-model cells are computed once and reused for every checkpoint.

Env: ADAPTER_REPO (required), BASE_MODEL, N_TASKS (default 100), POLL_S (600).
Stops when eval/final_results.json appears (the training job's own final eval).
"""

import json
import os
import subprocess
import time
import urllib.request
from pathlib import Path

from huggingface_hub import HfApi, snapshot_download

REPO = os.environ["ADAPTER_REPO"]
BASE_MODEL = os.environ.get("BASE_MODEL", "LiquidAI/LFM2.5-2.6B")
POLL_S = int(os.environ.get("POLL_S", "600"))
VLLM = "http://127.0.0.1:8001"

api = HfApi()


def repo_files():
    return set(api.list_repo_files(REPO))


def pending_checkpoints(fs):
    ckpts = sorted(
        {f.split("/")[0] for f in fs if f.startswith("checkpoint-") and f.endswith("adapter_config.json")},
        key=lambda s: int(s.rsplit("-", 1)[1]),
    )
    done = {f[len("eval/"):-len("_results.json")] for f in fs if f.startswith("eval/") and f.endswith("_results.json")}
    return [c for c in ckpts if c not in done]


def run_eval(model_name, out_path):
    env = dict(os.environ, ONLY_MODEL=model_name, OUT=out_path)
    subprocess.run(["python3", "/app/eval_adapter.py"], env=env, check=True)


def upload(local, remote):
    api.upload_file(path_or_fileobj=str(local), path_in_repo=remote, repo_id=REPO)


def lora_rpc(route, name, path=None):
    payload = {"lora_name": name}
    if path is not None:
        payload["lora_path"] = str(path)
    req = urllib.request.Request(
        f"{VLLM}/v1/{route}",
        data=json.dumps(payload).encode(),
        headers={"Content-Type": "application/json"},
    )
    urllib.request.urlopen(req, timeout=900)


def agg(cells_path):
    """pass@1 plus efficiency means over solved tasks (the article's second axis)."""
    cells = json.loads(Path(cells_path).read_text())["cells"]
    scored = [c["correct"] for c in cells.values() if not c["failed"]]
    solved = [c["metrics"] for c in cells.values() if c["correct"] and c.get("metrics")]

    def mean(key):
        vals = [m[key] for m in solved if key in m]
        return round(sum(vals) / len(vals), 1) if vals else None

    return {
        "pass_at_1": (sum(scored) / len(scored)) if scored else None,
        "scorable": len(scored),
        "solved": len(solved),
        "tool_calls_solved_mean": mean("tool_calls"),
        "turns_solved_mean": mean("turns"),
        "completion_tokens_solved_mean": mean("completion_tokens"),
        "prompt_tokens_solved_mean": mean("prompt_tokens_total"),
    }


def main():
    fs = repo_files()
    # 1. Base-model cells, once. Every checkpoint compares against these.
    #    Recomputed if the cached file predates per-cell metrics.
    if "eval/base_cells.json" in fs:
        local = api.hf_hub_download(REPO, "eval/base_cells.json")
        Path("/tmp/base_cells.json").write_text(Path(local).read_text())
    if "eval/base_cells.json" not in fs or any(
        c.get("metrics") is None
        for c in json.loads(Path("/tmp/base_cells.json").read_text())["cells"].values()
        if c["correct"]
    ):
        print("[watch] evaluating base model first", flush=True)
        run_eval(BASE_MODEL, "/tmp/base_cells.json")
        upload("/tmp/base_cells.json", "eval/base_cells.json")

    summary = {}
    if "eval/summary.json" in fs:
        summary = json.loads(Path(api.hf_hub_download(REPO, "eval/summary.json")).read_text())

    loaded = []
    while True:
        fs = repo_files()
        if "eval/final_results.json" in fs:
            print("[watch] training job's final eval landed; stopping", flush=True)
            return
        for ckpt in pending_checkpoints(fs):
            print(f"[watch] new checkpoint: {ckpt}", flush=True)
            local = Path(snapshot_download(REPO, allow_patterns=[f"{ckpt}/*"])) / ckpt
            lora_rpc("load_lora_adapter", ckpt, local)
            loaded.append(ckpt)
            if len(loaded) > 4:  # stay under --max-loras
                lora_rpc("unload_lora_adapter", loaded.pop(0))
            try:
                run_eval(ckpt, f"/tmp/{ckpt}_cells.json")
            except subprocess.CalledProcessError as exc:
                # Transient eval failure (client reconnect flake, sandbox hiccup):
                # skip it — the next poll recomputes pending and retries it.
                print(f"[watch] {ckpt} eval failed ({exc}); will retry next poll", flush=True)
                continue
            upload(f"/tmp/{ckpt}_cells.json", f"eval/{ckpt}_results.json")
            a, b = agg(f"/tmp/{ckpt}_cells.json"), agg("/tmp/base_cells.json")
            entry = dict(a)
            entry["base_pass_at_1"] = b["pass_at_1"]
            if a["pass_at_1"] is not None and b["pass_at_1"] is not None:
                entry["delta"] = a["pass_at_1"] - b["pass_at_1"]
            for key in ("tool_calls_solved_mean", "turns_solved_mean",
                        "completion_tokens_solved_mean", "prompt_tokens_solved_mean"):
                if a.get(key) is not None and b.get(key):
                    entry[key.replace("_mean", "_vs_base")] = round(a[key] / b[key], 3)
            summary[ckpt] = entry
            Path("/tmp/summary.json").write_text(json.dumps(summary, indent=2))
            upload("/tmp/summary.json", "eval/summary.json")
            print(f"[watch] {ckpt}: pass@1={a['pass_at_1']} (base {b['pass_at_1']}), "
                  f"tool_calls={a['tool_calls_solved_mean']} (base {b['tool_calls_solved_mean']})",
                  flush=True)
        time.sleep(POLL_S)


if __name__ == "__main__":
    main()