"""Download pinned source and create independent environments. Never trains/publishes."""
import argparse
import json
import os
from pathlib import Path
import shutil
import subprocess
import urllib.request

from common.runtime import ROOT, VENDOR, write, sha

FORMA = "3469f527c5f6833d654f62653ee92e0f73854b33"
FLEET = "0cc581616f7e5493e2e3999e71c134cedc879a6d"
V2 = "1648582b07a0f7f9d482a8c0f4cf89e6ecd86972"
CODE = "ee2ffc67fe0f9537d4928e158440c7d228a63100"

def get(url, dest):
    if dest.exists():
        return
    dest.parent.mkdir(parents=True, exist_ok=True)
    req = urllib.request.Request(url, headers={"User-Agent": "m4-local-training/1"})
    with urllib.request.urlopen(req, timeout=120) as r, dest.with_suffix(dest.suffix+".part").open("wb") as f:
        shutil.copyfileobj(r, f)
    dest.with_suffix(dest.suffix+".part").replace(dest)

def snapshot(repo, rev, folder):
    url = f"https://huggingface.co/api/models/{repo}/revision/{rev}"
    with urllib.request.urlopen(url, timeout=60) as r:
        metadata = json.load(r)
    for item in metadata["siblings"]:
        name = item["rfilename"]
        path = Path(name)
        if path.is_absolute() or ".." in path.parts:
            raise ValueError("Invalid remote path")
        if path.suffix in (".pt", ".cact", ".safetensors"):
            continue
        get(f"https://huggingface.co/{repo}/resolve/{rev}/{name}", folder / name)
    write(folder / "PIN.json", {"repo": repo, "revision": rev})

def run(*args, **kw):
    print("+", " ".join(map(str,args)), flush=True)
    subprocess.run(list(map(str,args)), check=True, **kw)

def env(name, requirement):
    dest = ROOT / ".venvs" / name
    lock = ROOT / "locks" / (name + ".txt")
    if not (dest / "bin/python").exists():
        run("uv", "venv", "--python", "3.12", dest)
    args = ["uv", "pip", "install", "--link-mode", "copy", "--python", dest / "bin/python", "-r", lock if lock.exists() else requirement]
    if not (dest / ".setup-complete").exists():
        args.append("--reinstall")
    run(*args)
    (dest / ".setup-complete").touch()
    lock.parent.mkdir(exist_ok=True)
    lock.write_bytes(subprocess.check_output(["uv", "pip", "freeze", "--python", str(dest / "bin/python")]))

def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("model", choices=["forma", "fleet", "all"])
    args = parser.parse_args()
    if not shutil.which("uv"):
        raise SystemExit("Install uv first: brew install uv")
    os.environ.setdefault("UV_CACHE_DIR", str(ROOT / ".cache/uv"))
    if args.model in ("forma", "all"):
        snapshot("Lucasgitguy-co/forma3d-tiny-3", FORMA, VENDOR / "forma")
        get(f"https://huggingface.co/Lucasgitguy-co/forma3d-tiny-2/resolve/{V2}/scripts/prepare_hf.py", VENDOR / "forma/prepare_hf.py")
        env("forma", VENDOR / "forma/recipe/requirements.txt")
        run("npm", "ci", "--ignore-scripts", cwd=VENDOR / "forma/recipe")
    if args.model in ("fleet", "all"):
        snapshot("Lucasgitguy-co/Fleet-1", FLEET, VENDOR / "fleet")
        env("fleet", VENDOR / "fleet/requirements.lock")
        # The runtime explicitly pins JAX 0.4.38 for its Metal extra.
        # An isolated older JAX/Flax stack avoids corrupting the reproduction environment.
        metal = ROOT / "fleet/metal-requirements.txt"
        try:
            env("fleet-metal", metal)
            write(ROOT / "locks/metal-setup.json", {"installed": True})
        except subprocess.CalledProcessError as exc:
            write(ROOT / "locks/metal-setup.json", {"installed": False, "error": str(exc)})
            print("Metal environment unavailable; the runner will use the CPU environment.", flush=True)
    files = {str(p.relative_to(VENDOR)): sha(p) for p in VENDOR.rglob("*") if p.is_file() and "node_modules" not in p.parts}
    write(ROOT / "locks/source-hashes.json", files)
    print("Setup complete. No weights trained or uploaded.")

if __name__ == "__main__":
    main()