from __future__ import annotations

import argparse
import json
import os
import sys
import hashlib
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

def main():
    ap = argparse.ArgumentParser(description="Run official Wan-Animate-2 with packed OrbitQuant W4A4.")
    ap.add_argument("--official-repo", default="/content/Wan-Animate-2")
    ap.add_argument("--model-root", default="/content/Wan2_2_Animate_Repo")
    ap.add_argument("--packed", default="/content/OrbitQuant_Animate2_W4A4_PACKED")
    ap.add_argument("--image", default="/content/wan_reference/reference.jpg")
    ap.add_argument("--video", default="/content/oAfghkYL_720p.mp4")
    ap.add_argument("--output", default="/content/OrbitQuant_Animate2_W4A4_PACKED/output_smoke")
    ap.add_argument("--prompt", default="A realistic person matching the reference image, performing the actions from the driving video, natural motion, detailed, high quality.")
    ap.add_argument("--prompt-ref", default="人物动作的参考视频")
    ap.add_argument("--negative-prompt", default="")
    ap.add_argument("--width", type=int, default=256)
    ap.add_argument("--height", type=int, default=320)
    ap.add_argument("--fps", type=int, default=24)
    ap.add_argument("--clip-len", type=int, default=17)
    ap.add_argument("--steps", type=int, default=10)
    ap.add_argument("--guidance", type=float, default=1.0)
    ap.add_argument("--seed", type=int, default=42)
    ap.add_argument("--placement", choices=["auto", "resident", "stream"], default="auto")
    ap.add_argument("--attention", choices=["official", "sol", "para", "hybrid"], default="hybrid")
    ap.add_argument("--sol-tau", type=float, default=1.0)
    ap.add_argument("--kv-cache", choices=["cpu", "cpu-pinned", "gpu"], default="cpu")
    ap.add_argument("--max-frames", type=int, default=17, help="Smoke input cap at target FPS; use 0 for the full driving video")
    ap.add_argument("--kernel-gate", default="/content/orbitquant_w4a4_cuda_gate.json")
    ap.add_argument("--skip-kernel-gate", action="store_true")
    args = ap.parse_args()

    if not args.skip_kernel_gate:
        gate_path = Path(args.kernel_gate)
        if not gate_path.is_file():
            raise RuntimeError(
                f"CUDA W4A4 kernel gate not found: {gate_path}. "
                "Run scripts/kernel_selftest.py successfully before loading the 14B model."
            )
        gate = json.loads(gate_path.read_text())
        if gate.get("status") != "PASS":
            raise RuntimeError(f"CUDA W4A4 kernel gate is not PASS: {gate}")
        import torch
        if gate.get("gpu") != torch.cuda.get_device_name():
            raise RuntimeError(
                f"kernel gate was produced on {gate.get('gpu')!r}, current GPU is {torch.cuda.get_device_name()!r}"
            )
        if gate.get("torch") != torch.__version__ or gate.get("cuda") != torch.version.cuda:
            raise RuntimeError(
                "Torch/CUDA changed since the kernel gate; rerun scripts/kernel_selftest.py "
                f"(gate torch={gate.get('torch')} cuda={gate.get('cuda')}, "
                f"current torch={torch.__version__} cuda={torch.version.cuda})."
            )
        artifact_gate = gate.get("packed_artifact")
        if not artifact_gate:
            raise RuntimeError(
                "CUDA gate did not validate the actual packed artifact. "
                "Rerun scripts/kernel_selftest.py --packed-dir <packed-dir>."
            )
        pdir_gate = Path(args.packed)
        manifest_path = pdir_gate / "packed_manifest.json"
        if not manifest_path.is_file():
            raise RuntimeError(f"packed manifest missing: {manifest_path}")
        current_manifest_sha = hashlib.sha256(manifest_path.read_bytes()).hexdigest()
        if current_manifest_sha != artifact_gate.get("manifest_sha256"):
            raise RuntimeError("packed manifest changed since the CUDA artifact gate")
        for name, size in artifact_gate.get("shard_sizes", {}).items():
            p = pdir_gate / name
            if not p.is_file() or p.stat().st_size != int(size):
                raise RuntimeError(f"packed shard changed since CUDA gate: {p}")
        print("✓ CUDA W4A4 kernel + actual packed-artifact gate:", gate_path)

    official = Path(args.official_repo).resolve()
    if not official.is_dir():
        raise FileNotFoundError(f"official Wan-Animate-2 source not found: {official}; run COLAB_BOOTSTRAP.py first")
    sys.path.insert(0, str(official))

    from orbitquant_wan_a2.official_runtime import run_official_packed_w4a4

    report = run_official_packed_w4a4(
        official_repo=official,
        model_root=args.model_root,
        packed_dir=args.packed,
        reference_image=args.image,
        driving_video=args.video,
        output_dir=args.output,
        prompt=args.prompt,
        prompt_ref=args.prompt_ref,
        negative_prompt=args.negative_prompt,
        width=args.width,
        height=args.height,
        fps=args.fps,
        clip_len=args.clip_len,
        steps=args.steps,
        guidance_scale=args.guidance,
        seed=args.seed,
        transformer_placement=args.placement,
        attention=args.attention,
        sol_tau=args.sol_tau,
        kv_cache_placement=args.kv_cache,
        max_input_frames=(None if args.max_frames <= 0 else args.max_frames),
    )
    if int(os.environ.get("RANK", "0")) == 0:
        print(json.dumps(report, indent=2, default=str))


if __name__ == "__main__":
    main()