# /// script
# requires-python = ">=3.11"
# dependencies = [
#     "torch>=2.4",
#     "transformers>=4.45",
#     "unsloth",
#     "peft>=0.13",
#     "bitsandbytes>=0.43",
#     "accelerate>=0.34",
#     "fastapi>=0.111",
#     "pydantic>=2.7",
#     "numpy>=1.26",
#     "huggingface-hub>=0.24",
# ]
# ///
"""UV-script wrapper that runs Frozen + Trained Qwen evals on T4.

Clones the env code, downloads the trained LoRA checkpoint from the runs
dataset, runs 50-episode evals against the local in-process MCPArenaEnv
on phase=eval for both the frozen base model AND the trained LoRA-merged
model, then uploads the eval JSONLs back to the runs dataset.

Usage from local machine:
  hf jobs uv run --flavor t4-medium -s HF_TOKEN \
    --env CHECKPOINT_PATH=phase_1/20260425-191159/grpo_phase_1/final \
    --env N_EPISODES=50 \
    scripts/launch_eval_t4.py
"""
import json
import os
import subprocess
import sys
from datetime import datetime
from pathlib import Path


def run(cmd: list[str], **kw) -> None:
    print(f"[launch_eval] $ {' '.join(cmd)}", flush=True)
    subprocess.run(cmd, check=True, **kw)


def main() -> None:
    n_episodes = int(os.environ.get("N_EPISODES", "50"))
    repo_url = os.environ.get("REPO_URL", "https://huggingface.co/spaces/vex-0/mcparena")
    runs_dataset = os.environ.get("RUNS_DATASET", "vex-0/mcparena-runs")
    checkpoint_subpath = os.environ.get("CHECKPOINT_PATH", "phase_1/20260425-191159/grpo_phase_1/final")
    # On a re-eval (e.g. Phase 2) we already have a Frozen Qwen baseline; skip it.
    skip_frozen = os.environ.get("SKIP_FROZEN", "").lower() in ("1", "true", "yes")
    # Tag the trained eval log filename so re-runs for different checkpoints don't overwrite each other.
    trained_label = os.environ.get("TRAINED_LABEL", "trained_qwen")
    frozen_label = os.environ.get("FROZEN_LABEL", "frozen_qwen")
    # Eval phase: "eval" (default, held-out) or "phase_2" / "phase_1" for in-distribution checks.
    eval_phase = os.environ.get("EVAL_PHASE", "eval")
    eval_temperature = float(os.environ.get("EVAL_TEMP", "0.7"))
    workdir = "/tmp/mcparena"

    # Clone env code from HF Space (master branch is named 'main' on Spaces)
    run(["git", "clone", repo_url, workdir])
    os.chdir(workdir)

    # PYTHONPATH=workdir makes server.* importable without pip install -e .
    # PYTHONUNBUFFERED ensures print() output streams immediately in HF Job logs.
    env = dict(os.environ, PYTHONPATH=workdir, PYTHONUNBUFFERED="1")

    # Download the trained LoRA checkpoint from the dataset
    from huggingface_hub import snapshot_download
    ckpt_dir = snapshot_download(
        repo_id=runs_dataset,
        repo_type="dataset",
        allow_patterns=f"{checkpoint_subpath}/**",
        local_dir="/tmp/dl",
    )
    # The actual checkpoint lands at /tmp/dl/<checkpoint_subpath>/
    full_ckpt = os.path.join("/tmp/dl", checkpoint_subpath)
    print(f"[launch_eval] checkpoint at: {full_ckpt}", flush=True)
    print(
        f"[launch_eval] checkpoint contents: "
        f"{os.listdir(full_ckpt) if os.path.isdir(full_ckpt) else 'NOT A DIR'}",
        flush=True,
    )

    out_dir = Path("runs/eval")
    out_dir.mkdir(parents=True, exist_ok=True)

    eval_failed = False

    # --- Frozen Qwen baseline ---
    if skip_frozen:
        print("[launch_eval] SKIP_FROZEN=1 — reusing prior frozen_qwen_eval.jsonl", flush=True)
    else:
        print(f"[launch_eval] running Frozen Qwen eval ({n_episodes} eps on phase=eval)", flush=True)
        try:
            run(
                [
                    sys.executable,
                    "-c",
                    _make_eval_command(
                        agent=frozen_label,
                        n_episodes=n_episodes,
                        out_path=str(out_dir / f"{frozen_label}_eval.jsonl"),
                        checkpoint_path=None,
                        eval_phase=eval_phase,
                        eval_temperature=eval_temperature,
                    ),
                ],
                env=env,
            )
        except subprocess.CalledProcessError as e:
            print(f"[launch_eval] Frozen Qwen eval failed (exit {e.returncode}); continuing", flush=True)
            eval_failed = True

    # --- Trained Qwen (LoRA at checkpoint_subpath) ---
    print(f"[launch_eval] running Trained Qwen eval ({n_episodes} eps on phase=eval)", flush=True)
    try:
        run(
            [
                sys.executable,
                "-c",
                _make_eval_command(
                    agent=trained_label,
                    n_episodes=n_episodes,
                    out_path=str(out_dir / f"{trained_label}_eval.jsonl"),
                    checkpoint_path=full_ckpt,
                    eval_phase=eval_phase,
                    eval_temperature=eval_temperature,
                ),
            ],
            env=env,
        )
    except subprocess.CalledProcessError as e:
        print(f"[launch_eval] Trained Qwen eval failed (exit {e.returncode}); continuing", flush=True)
        eval_failed = True

    # Upload runs/eval/ to dataset regardless of eval success
    print("[launch_eval] uploading eval results to dataset", flush=True)
    from huggingface_hub import HfApi
    api = HfApi()
    ts = datetime.utcnow().strftime("%Y%m%d-%H%M%S")
    upload_path = f"eval/{ts}"
    api.upload_folder(
        folder_path=str(out_dir),
        repo_id=runs_dataset,
        repo_type="dataset",
        path_in_repo=upload_path,
    )
    print(f"[launch_eval] uploaded → {runs_dataset}/{upload_path}", flush=True)

    if eval_failed:
        sys.exit(1)


def _make_eval_command(
    agent: str,
    n_episodes: int,
    out_path: str,
    checkpoint_path: str | None,
    eval_phase: str = "eval",
    eval_temperature: float = 0.7,
) -> str:
    """Return a self-contained Python source string for `python -c`."""
    return f"""
import json, os, sys
sys.path.insert(0, '/tmp/mcparena')
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel
from server.env import MCPArenaEnv
from server.prompts import SYSTEM_PROMPT
from server.reward.rubrics import compute_episode_breakdown


# Determine base model: prefer adapter_config.json's base_model_name_or_path
# (so 7B/14B LoRAs auto-pair with their training base), fall back to env var,
# then to default 3B.
ckpt_path = {repr(checkpoint_path)}
base_model = 'Qwen/Qwen2.5-3B-Instruct'
if ckpt_path:
    cfg_path = os.path.join(ckpt_path, 'adapter_config.json')
    if os.path.exists(cfg_path):
        try:
            with open(cfg_path) as f:
                cfg = json.load(f)
            if cfg.get('base_model_name_or_path'):
                base_model = cfg['base_model_name_or_path']
        except Exception as exc:
            print(f'[eval] WARN reading adapter_config: {{exc}}', flush=True)
base_model = os.environ.get('EVAL_BASE_MODEL', base_model)
print(f'[eval] base model: {{base_model}}', flush=True)

tokenizer = AutoTokenizer.from_pretrained(base_model)
base = AutoModelForCausalLM.from_pretrained(
    base_model,
    torch_dtype='auto',
    device_map='cuda',
)

if ckpt_path:
    print(f'[eval] loading LoRA adapter from {{ckpt_path}}', flush=True)
    model = PeftModel.from_pretrained(base, ckpt_path).eval()
else:
    print('[eval] using frozen base model', flush=True)
    model = base.eval()


def render_user_prompt(observation):
    catalog = observation.get('catalog', [])
    catalog_text = '\\n'.join(
        f'- {{t["name"]}}: {{t["description"]}} (cost: {{t["cost_per_call"]}})'
        for t in catalog[:60]
    )
    return (
        f'Task: {{observation["task_text"]}}\\n\\n'
        f'Catalog (showing {{len(catalog)}} tools):\\n{{catalog_text}}\\n\\n'
        f'Budget remaining: {{observation["budget_remaining"]}}\\n'
        f'Step: {{observation["step"]}}/{{observation["step_cap"]}}\\n\\n'
        f'Emit ONE action as JSON: {{{{thought, action, confidence}}}}.'
    )


def run_one_episode(seed):
    env = MCPArenaEnv()
    obs = env.reset(seed=seed, phase={repr(eval_phase)})
    done = False
    while not done and env.state.step < 12:
        chat = [
            {{'role': 'system', 'content': SYSTEM_PROMPT}},
            {{'role': 'user', 'content': render_user_prompt(obs)}},
        ]
        prompt_text = tokenizer.apply_chat_template(
            chat, tokenize=False, add_generation_prompt=True
        )
        toks = tokenizer(prompt_text, return_tensors='pt').to(model.device)
        out_ids = model.generate(
            **toks,
            max_new_tokens=96,
            do_sample=True,
            temperature={eval_temperature},
            top_p=0.9,
            pad_token_id=tokenizer.eos_token_id,
        )
        response_text = tokenizer.decode(
            out_ids[0][toks['input_ids'].shape[1]:], skip_special_tokens=True
        )
        out = env.step(response_text)
        obs = out['observation']
        done = out['done']
    breakdown = compute_episode_breakdown(env.state)
    return {{
        'seed': seed,
        'phase': {repr(eval_phase)},
        'agent': '{agent}',
        'final_reward': breakdown['total'],
        'steps': env.state.step,
        'breakdown': breakdown,
    }}


with open({repr(out_path)}, 'w') as f:
    for ep_idx in range({n_episodes}):
        seed = ep_idx * 31 + 7
        try:
            rec = run_one_episode(seed)
            f.write(json.dumps(rec) + '\\n')
            f.flush()
            if ep_idx % 5 == 0:
                print(
                    f'[eval-{agent}] ep {{ep_idx}}: '
                    f'reward={{rec["final_reward"]:.3f}} '
                    f'task={{rec["breakdown"]["raw"]["task_success"]:.2f}}',
                    flush=True,
                )
        except Exception as e:
            print(f'[eval-{agent}] ep {{ep_idx}} crashed: {{e}}; skipping', flush=True)

print(f'[eval-{agent}] done; wrote to {out_path}', flush=True)
"""


if __name__ == "__main__":
    main()