from __future__ import annotations

import json
import os
from pathlib import Path
import shutil
import subprocess
import sys
import site

ROOT = Path(__file__).resolve().parents[1]
PINS = json.loads((ROOT / "source_pins.json").read_text())
CONTENT = Path("/content")


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


def clone_at(name: str, destination: Path):
    spec = PINS[name]
    sha = spec["commit"]
    if destination.exists():
        shutil.rmtree(destination)
    run(["git", "init", "-q", destination])
    run(["git", "-C", destination, "remote", "add", "origin", spec["url"]])
    run(["git", "-C", destination, "fetch", "-q", "--depth=1", "origin", sha])
    run(["git", "-C", destination, "checkout", "-q", "FETCH_HEAD"])
    got = subprocess.check_output(["git", "-C", destination, "rev-parse", "HEAD"], text=True).strip()
    if got != sha:
        raise RuntimeError(f"{name} pin mismatch: expected {sha}, got {got}")
    print(f"✓ {name} {got}")


def version_tuple(text: str):
    base = text.split("+")[0].split("rc")[0]
    vals = []
    for x in base.split(".")[:3]:
        try:
            vals.append(int(x))
        except Exception:
            vals.append(0)
    return tuple(vals + [0] * (3 - len(vals)))


def main():
    import torch

    print("Torch:", torch.__version__)
    print("CUDA :", torch.version.cuda)
    if version_tuple(torch.__version__) < (2, 10, 0):
        raise RuntimeError(
            "Sol-Attn's pinned sol-engine backend requires PyTorch >= 2.10. "
            "Use a Colab runtime with Torch 2.10+ rather than silently replacing the CUDA stack here."
        )
    if not torch.cuda.is_available():
        raise RuntimeError("CUDA GPU is required")
    if torch.version.cuda is None or version_tuple(torch.version.cuda) < (12, 8, 0):
        raise RuntimeError(
            f"Pinned Sol-Attn requires CUDA >= 12.8; current Torch CUDA runtime is {torch.version.cuda!r}. "
            "Use a newer Colab GPU runtime instead of silently replacing PyTorch/CUDA."
        )
    cc = torch.cuda.get_device_capability()
    print("GPU  :", torch.cuda.get_device_name())
    print("CC   :", cc)
    if cc[0] < 8:
        raise RuntimeError("Sol-Attn supports NVIDIA SM80+ in this package")

    # Install runtime Python dependencies without replacing PyTorch.
    deps = [
        "triton>=3.6", "safetensors==0.7.0", "numpy>=1.26", "PyYAML", "easydict",
        "huggingface_hub==0.36.0", "transformers==4.57.6", "accelerate==1.13.0",
        "diffusers==0.36.0", "imageio", "imageio-ffmpeg", "decord", "opencv-python-headless",
        "moviepy", "loguru", "einops", "tqdm", "psutil", "omegaconf", "addict", "ftfy",
        "scipy", "sentencepiece", "protobuf", "packaging", "pillow"
    ]
    run([sys.executable, "-m", "pip", "install", "-q", "-U", *deps])

    wan = CONTENT / "Wan-Animate-2"
    sana = CONTENT / "Sana-sol-engine"
    para = CONTENT / "ParaAttention"
    clone_at("Wan-Video/Wan-Animate-2", wan)
    clone_at("NVlabs/Sana", sana)
    clone_at("chengzeyi/ParaAttention", para)

    # Verify the exact upstream Animate-2 distilled contract we patch into.
    import yaml
    source_cfg = yaml.safe_load((wan / "infer" / "wan_animate_2_distillation.yaml").read_text())
    tr = source_cfg["model"]["transformer"]
    expected = {
        "in_dim": 36,
        "dim": 5120,
        "ffn_dim": 13824,
        "num_heads": 40,
        "num_layers": 40,
        "log_scale": -1.3,
    }
    for key, value in expected.items():
        if tr.get(key) != value:
            raise RuntimeError(f"official Animate-2 source contract changed: {key}={tr.get(key)!r}, expected {value!r}")
    if source_cfg["model"].get("flow_solver") != "euler" or float(source_cfg["test_cfg"].get("sample_shift")) != 5.0:
        raise RuntimeError("official distilled Euler/shift-5 contract changed")
    print("✓ official Animate-2 architecture + distilled Euler contract")

    # These integrations are pure Python/Triton at runtime.  Put the exact
    # checked-out source directories on sys.path via one .pth file instead of
    # invoking setuptools-scm/editable builds that can drift or pull deps.
    paths = [
        str(ROOT),
        str(wan),
        str(sana / "techniques" / "sparse_backends"),
        str(para / "src"),
    ]
    site_dir = Path(site.getsitepackages()[0])
    pth = site_dir / "orbitquant_wan_a2_third_party.pth"
    pth.write_text("\n".join(paths) + "\n")
    for path in reversed(paths):
        if path not in sys.path:
            sys.path.insert(0, path)

    # Import gates catch namespace/path mistakes immediately.
    import orbitquant_wan_a2  # noqa: F401
    from sol_attn import sol_attn  # noqa: F401
    import para_attn.primitives  # noqa: F401
    from wanxiang.models.wan_animate_2_model import WanAnimate2Transformer  # noqa: F401

    print("\nBOOTSTRAP PASS")
    print("  official source :", wan)
    print("  Sol-Attn        :", sana / "techniques" / "sparse_backends")
    print("  ParaAttention   :", para)
    print("  runtime package :", ROOT)
    print("  path file       :", pth)
    print("\nNext: python", ROOT / "scripts" / "kernel_selftest.py")


if __name__ == "__main__":
    main()