# /// script
# requires-python = ">=3.11"
# dependencies = [
#     "torch>=2.4",
#     "transformers>=4.45",
#     "trl>=0.11",
#     "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 MCPArena GRPO training on an HF Jobs T4.

Clones the latest env from the HF Space, installs it in editable mode, then
invokes training/train_grpo.py. Pass phase + episodes via env vars:
  PHASE=phase_1 EPISODES=200

Usage from local machine:
  hf jobs uv run --flavor t4-medium -s HF_TOKEN \
    --env PHASE=phase_1 --env EPISODES=200 \
    scripts/launch_train_t4.py
"""
import os
import subprocess
import sys


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


def main() -> None:
    phase = os.environ.get("PHASE", "phase_1")
    episodes = int(os.environ.get("EPISODES", "200"))
    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")
    # Optional path inside the runs dataset to a prior LoRA adapter
    # (e.g. "phase_1/20260425-191159/grpo_phase_1/final" for Phase 2 warm start).
    resume_from_dataset_path = os.environ.get("RESUME_FROM_DATASET_PATH", "")
    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)

    # 'server' / 'client' / 'training' are importable from CWD; UV-managed venvs
    # don't ship pip, so we skip `pip install -e .` and rely on PYTHONPATH=cwd.
    # PYTHONUNBUFFERED ensures print() output streams immediately in HF Job logs.
    env = dict(os.environ, PYTHONPATH=workdir, PYTHONUNBUFFERED="1")

    extra_args: list[str] = []
    # REINFORCE baseline tuning — bump above the no-commit reward floor (~0.50) so
    # don't-commit episodes get negative gradient and the policy is pushed off the trap.
    reward_baseline = os.environ.get("REWARD_BASELINE", "")
    if reward_baseline:
        extra_args.extend(["--reward_baseline", reward_baseline])
    base_model = os.environ.get("BASE_MODEL", "")
    if base_model:
        extra_args.extend(["--base_model", base_model])
    if resume_from_dataset_path:
        from huggingface_hub import snapshot_download
        local_resume = "/tmp/resume_adapter"
        print(f"[launch] downloading resume adapter from {runs_dataset}/{resume_from_dataset_path}", flush=True)
        snapshot_download(
            repo_id=runs_dataset,
            repo_type="dataset",
            allow_patterns=[f"{resume_from_dataset_path}/*"],
            local_dir=local_resume,
        )
        resume_dir = os.path.join(local_resume, resume_from_dataset_path)
        if not os.path.isdir(resume_dir):
            raise SystemExit(f"[launch] expected resume dir {resume_dir} not found")
        print(f"[launch] resume adapter at {resume_dir}", flush=True)
        extra_args.extend(["--resume_from", resume_dir])

    train_failed = False
    try:
        run(
            [sys.executable, "training/train_grpo.py", "--phase", phase, "--episodes", str(episodes), *extra_args],
            env=env,
        )
    except subprocess.CalledProcessError as e:
        print(f"[launch] training failed with exit {e.returncode}; uploading partial runs/ anyway", flush=True)
        train_failed = True

    # Upload runs/ to the HF dataset regardless of training success
    runs_dir = os.path.join(workdir, "runs")
    if os.path.isdir(runs_dir):
        from huggingface_hub import HfApi, create_repo
        api = HfApi()
        try:
            create_repo(runs_dataset, repo_type="dataset", exist_ok=True, private=True)
        except Exception as e:
            print(f"[launch] create_repo warning (likely already exists): {e}", flush=True)
        # Upload with phase-prefixed path so multi-phase runs don't collide
        from datetime import datetime
        ts = datetime.utcnow().strftime("%Y%m%d-%H%M%S")
        path_in_repo = f"{phase}/{ts}"
        api.upload_folder(
            folder_path=runs_dir,
            repo_id=runs_dataset,
            repo_type="dataset",
            path_in_repo=path_in_repo,
        )
        print(f"[launch] uploaded {runs_dir} → dataset {runs_dataset}/{path_in_repo}", flush=True)
    else:
        print(f"[launch] no runs/ directory to upload", flush=True)

    if train_failed:
        sys.exit(1)


if __name__ == "__main__":
    main()