"""Held-out evaluation: base model vs LoRA adapter on SmolDataEnvs test tasks.

Runs N test tasks through the env server under OpenCode, once per model, and
reports pass@1 (first graded attempt correct) for each, plus the per-task
agreement breakdown the article's comparisons used.

The vLLM server must serve the base model AND the adapter as a named LoRA:
    vllm serve LiquidAI/LFM2.5-2.6B --enable-lora --max-lora-rank 32 \
        --lora-modules trained=<adapter-repo> ...

Config via env:
    ENV_URL        env server (default http://127.0.0.1:8000)
    VLLM_URL       policy server (default http://127.0.0.1:8001)
    TEST_SPLIT     served test split (default /app/prepared/datasets/test)
    BASE_MODEL     served base model id (default LiquidAI/LFM2.5-2.6B)
    ADAPTER_NAME   LoRA name on the server (default "trained"); "" evaluates base only
    N_TASKS        how many test tasks (default 50)
    CONCURRENCY    parallel rollouts (default 8; keep <= MAX_CONCURRENT_ENVS)
    OUT            results JSON path (default /app/eval_results.json)
    ONLY_MODEL     evaluate just this one served model name (per-checkpoint eval)
"""

import json
import os
from concurrent.futures import ThreadPoolExecutor, as_completed

from harbor_env import HarborEnv  # re-export of openenv.harbor.client.HarborEnv

ENV_URL = os.environ.get("ENV_URL", "http://127.0.0.1:8000")
VLLM_URL = os.environ.get("VLLM_URL", "http://127.0.0.1:8001")
TEST_SPLIT = os.environ.get("TEST_SPLIT", "/app/prepared/datasets/test")
BASE_MODEL = os.environ.get("BASE_MODEL", "LiquidAI/LFM2.5-2.6B")
ADAPTER_NAME = os.environ.get("ADAPTER_NAME", "trained")
ONLY_MODEL = os.environ.get("ONLY_MODEL", "")
N_TASKS = int(os.environ.get("N_TASKS", "50"))
CONCURRENCY = int(os.environ.get("CONCURRENCY", "8"))
OUT = os.environ.get("OUT", "/app/eval_results.json")

def run_one(task_index, model):
    """One graded rollout. Returns (correct, failed, error_or_none).

    A FRESH client per rollout: HarborSessionFactory documents that shared
    clients break under concurrency, and a reused websocket hits an openenv
    reconnect bug (mcp_client close() returning None) that kills the run.
    """
    env = HarborEnv(base_url=ENV_URL)
    kwargs = dict(
        split=TEST_SPLIT,
        task_index=task_index,
        harness="opencode",
        sandbox=os.environ.get("SANDBOX", "hf-sandbox"),
        reward_key="correctness,reward",
        llm_url=VLLM_URL,
        model=model,
        purpose="eval",   # required by the server when eval_sampling is passed
        eval_sampling={"temperature": 1.0, "top_p": 1.0},
        agent_step_limit=17,      # production eval cap from the FineEnvs Space
        agent_timeout_sec=600,
    )
    try:
        try:
            result = env.run_rollout(**kwargs)
        except TypeError:
            # older server without eval_sampling: retry without it
            kwargs.pop("eval_sampling")
            result = env.run_rollout(**kwargs)
    except Exception as exc:
        return task_index, model, None, True, f"{type(exc).__name__}: {exc}", None
    finally:
        try:
            env.close()
        except Exception:
            pass
    if not result.ok or result.reward is None:
        return task_index, model, None, True, getattr(result, "error", None) or "not ok", None
    turns = getattr(result, "turns", []) or []
    metrics = {
        "turns": getattr(result, "n_turns", 0) or len(turns),
        "tool_calls": sum(getattr(t, "n_tools", 0) for t in turns),
        "completion_tokens": sum(len(getattr(t, "completion_token_ids", []) or []) for t in turns),
        # Each turn resends the whole conversation, so this grows ~quadratically
        # with turn count — the quantity the article's input-token savings cut.
        "prompt_tokens_total": sum(len(getattr(t, "prompt_token_ids", []) or []) for t in turns),
        "wall_s": getattr(result, "wall_s", 0.0),
    }
    return task_index, model, float(result.reward) == 1.0, False, None, metrics


def main():
    if ONLY_MODEL:
        models = [ONLY_MODEL]
    else:
        models = [BASE_MODEL] + ([ADAPTER_NAME] if ADAPTER_NAME else [])
    work = [(i, m) for i in range(N_TASKS) for m in models]
    cells = {}  # (task_index, model) -> {"correct": bool|None, "failed": bool, "error": str|None}
    with ThreadPoolExecutor(max_workers=CONCURRENCY) as pool:
        futures = [pool.submit(run_one, i, m) for i, m in work]
        done = 0
        for fut in as_completed(futures):
            i, m, correct, failed, err, metrics = fut.result()
            cells[(i, m)] = {"correct": correct, "failed": failed, "error": err, "metrics": metrics}
            done += 1
            if done % 20 == 0 or done == len(work):
                print(f"[eval] {done}/{len(work)} cells done", flush=True)

    summary = {"n_tasks": N_TASKS, "models": {}}
    for m in models:
        scorable = [c["correct"] for (i, mm), c in cells.items() if mm == m and not c["failed"]]
        failed = sum(1 for (i, mm), c in cells.items() if mm == m and c["failed"])
        summary["models"][m] = {
            "pass_at_1": sum(scorable) / len(scorable) if scorable else None,
            "solved": sum(scorable),
            "scorable": len(scorable),
            "failed_rollouts": failed,
        }
    if len(models) == 2:
        b, a = models
        both = sum(1 for i in range(N_TASKS)
                   if cells.get((i, b), {}).get("correct") and cells.get((i, a), {}).get("correct"))
        only_a = sum(1 for i in range(N_TASKS)
                     if not cells.get((i, b), {}).get("correct") and cells.get((i, a), {}).get("correct"))
        only_b = sum(1 for i in range(N_TASKS)
                     if cells.get((i, b), {}).get("correct") and not cells.get((i, a), {}).get("correct"))
        summary["agreement"] = {"both_solved": both, "only_adapter": only_a, "only_base": only_b}
        pa = summary["models"][a]["pass_at_1"]
        pb = summary["models"][b]["pass_at_1"]
        if pa is not None and pb is not None:
            summary["delta"] = pa - pb

    print(json.dumps(summary, indent=2))
    with open(OUT, "w") as f:
        json.dump({"summary": summary,
                   "cells": {f"{i}|{m}": c for (i, m), c in cells.items()}}, f, indent=2)
    print(f"[eval] wrote {OUT}")


if __name__ == "__main__":
    main()