"""Plan/check an isolated pinned Lightning pilot; --execute launches real GPU RL."""

from __future__ import annotations

import argparse
import ast
import hashlib
import json
import os
from pathlib import Path
import secrets
import signal
import socket
import subprocess
import time
import urllib.error
import urllib.request


REVISION = "69f5a40b36d9bff646b8e7a5532e8b78bd41eaf9"
ROOT = Path(__file__).resolve().parents[2]


def check_source(upstream: Path) -> dict:
    revision = subprocess.check_output(["git", "-C", str(upstream), "rev-parse", "HEAD"], text=True).strip()
    if revision != REVISION:
        raise ValueError(f"Expected upstream {REVISION}; checkout has {revision}")
    if subprocess.check_output(["git", "-C", str(upstream), "status", "--porcelain", "--untracked-files=no"], text=True).strip():
        raise ValueError("Pinned upstream must have no tracked modifications")
    signatures = {
        "examples/calc_x/train_calc_agent.py": ("build_config", {"model", "agl_base_url", "agl_key", "config_overrides"}),
        "agentlightning/verl/entrypoint.py": ("run_ppo", {"config", "train_dataset", "val_dataset"}),
    }
    for path, (name, required) in signatures.items():
        tree = ast.parse((upstream / path).read_text())
        function = next((node for node in tree.body if isinstance(node, ast.FunctionDef) and node.name == name), None)
        arguments = set() if function is None else {arg.arg for arg in function.args.args + function.args.kwonlyargs}
        if not required <= arguments:
            raise ValueError(f"Upstream API contract changed: {path}:{name}")
    return {"revision": revision, "checked_functions": [value[0] for value in signatures.values()]}


def _stop(process: subprocess.Popen) -> None:
    if process.poll() is None:
        os.killpg(process.pid, signal.SIGTERM)
        try:
            process.wait(timeout=10)
        except subprocess.TimeoutExpired:
            os.killpg(process.pid, signal.SIGKILL)
            process.wait()


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--upstream", type=Path, required=True)
    parser.add_argument("--model", type=Path, help="Required for --execute: local immutable snapshot")
    parser.add_argument("--output", type=Path, default=Path("runs/agent-lightning"))
    parser.add_argument("--recipe", type=Path, default=ROOT / "recipes/agent-lightning.json")
    parser.add_argument("--execute", action="store_true")
    args = parser.parse_args()
    upstream = args.upstream.resolve()
    source = check_source(upstream)
    recipe = json.loads(args.recipe.read_text())
    if recipe["upstream_revision"] != REVISION:
        raise ValueError("Recipe and launcher upstream revisions must agree")
    gateway = f"http://127.0.0.1:{recipe['port']}"
    plan = {"status": "planned", "qualification": "source-checked; GPU run not qualified",
            "source": source, "model": str(args.model.resolve()) if args.model else None,
            "output": str(args.output.resolve()), "gateway": gateway,
            "overrides": recipe["overrides"], "seed": recipe["seed"]}
    if not args.execute:
        print(json.dumps(plan, indent=2))
        return
    if args.model is None or not args.model.is_dir():
        raise ValueError("--execute requires an existing local --model snapshot")
    python = upstream / ".venv/bin/python"
    if not python.is_file():
        raise ValueError("Prepare the isolated upstream .venv with uv sync and its pinned setup_verl.sh")
    output = args.output.resolve()
    with socket.socket() as probe:
        probe.bind(("127.0.0.1", recipe["port"]))
    output.mkdir(parents=True, exist_ok=False)
    env = dict(os.environ)
    env["PYTHONPATH"] = os.pathsep.join([str(upstream), str(Path(__file__).resolve().parent)])
    env.update(HF_HUB_OFFLINE="1", TRANSFORMERS_OFFLINE="1", WANDB_MODE="disabled",
               HESSIAN_AGL_KEY=secrets.token_hex(24))
    def digest(path: Path) -> str:
        with path.open("rb") as stream:
            return hashlib.file_digest(stream, "sha256").hexdigest()

    plan["model_files_sha256"] = {str(file.relative_to(args.model.resolve())): digest(file)
                                    for file in sorted(args.model.resolve().rglob("*")) if file.is_file()}
    plan["pilot_source_sha256"] = {file.name: digest(file)
                                    for file in Path(__file__).parent.glob("*.py")}
    plan["recipe_sha256"] = digest(args.recipe)
    processes, logs = [], []
    try:
        for binary, arguments in [
            ("agl-server", ["host=127.0.0.1", f"port={recipe['port']}", "key=" + env["HESSIAN_AGL_KEY"], "default_proxy.model_name=" + str(args.model.resolve())]),
            ("agl-controller", ["runner_type=local", "agl_server.url=" + gateway, "agl_server.key=" + env["HESSIAN_AGL_KEY"]]),
        ]:
            if binary == "agl-controller":
                for _ in range(60):
                    if processes[0].poll() is not None:
                        raise RuntimeError("Gateway exited; inspect agl-server.log")
                    try:
                        with urllib.request.urlopen(gateway + "/healthz", timeout=1):
                            break
                    except (urllib.error.URLError, TimeoutError):
                        time.sleep(0.5)
                else:
                    raise RuntimeError("Gateway did not become ready")
            log = (output / f"{binary}.log").open("w")
            logs.append(log)
            processes.append(subprocess.Popen([str(python.parent / binary), *arguments], cwd=upstream,
                                               env=env, stdout=log, stderr=subprocess.STDOUT,
                                               start_new_session=True))
        plan["status"] = "running"
        plan_path = output / "pilot-manifest.json"
        plan_path.write_text(json.dumps(plan, indent=2) + "\n")
        freeze = subprocess.check_output([str(python), "-m", "pip", "freeze"], env=env, text=True)
        (output / "environment.txt").write_text(freeze)
        subprocess.run([str(python), str(Path(__file__).with_name("train.py")), "--plan", str(plan_path)],
                       cwd=upstream, env=env, check=True)
        plan["status"] = "completed"
    except BaseException:
        plan["status"] = "failed"
        raise
    finally:
        for process in reversed(processes):
            _stop(process)
        for log in logs:
            log.close()
        (output / "pilot-manifest.json").write_text(json.dumps(plan, indent=2) + "\n")


if __name__ == "__main__":
    main()