"""Pinned ComfyUI core with checkpoint loading before GPU requests."""
from pathlib import Path
import subprocess
import sys
import gc

ROOT = Path(__file__).resolve().parent
COMFY_COMMIT = "65787d668397d230bf5839d69a0a7239e2dad378"
MODEL_ID = "Noctaluna/Noct-Q-Uncensored-Qwen-Image-2.1"
MODEL_REVISION = "a81b9af51120a78e285e57906f2250a2a02080e9"
COMPANION_ID = "Comfy-Org/Qwen-Image-2.1"
COMPANION_REVISION = "cb504a4090723e43f17ad01cec0359490e2de613"
CHECKPOINT = "NoctQ_V4_int8_convrot.safetensors"


def prepare_source() -> Path:
    """Fetch one pinned ComfyUI revision; never execute custom nodes."""
    target = ROOT / "ComfyUI"
    if not target.exists():
        target.mkdir()
        subprocess.run(["git", "init", str(target)], check=True)
        subprocess.run(["git", "-C", str(target), "remote", "add", "origin",
                        "https://github.com/Comfy-Org/ComfyUI.git"], check=True)
        subprocess.run(["git", "-C", str(target), "fetch", "--depth", "1",
                        "origin", COMFY_COMMIT], check=True)
        subprocess.run(["git", "-C", str(target), "checkout", "--detach",
                        "FETCH_HEAD"], check=True)
    actual = subprocess.check_output(
        ["git", "-C", str(target), "rev-parse", "HEAD"], text=True).strip()
    if actual != COMFY_COMMIT:
        raise RuntimeError("ComfyUI revision does not match the pinned runtime.")
    sys.path.insert(0, str(target))
    return target


def load_models():
    """Load the exact INT8 convrot files before the first GPU request."""
    # app.py must have imported spaces before calling this function.
    import torch
    from huggingface_hub import hf_hub_download

    prepare_source()
    import comfy.options
    comfy.options.enable_args_parsing()
    saved_argv = sys.argv[:]
    # Initialize Comfy's management module on CPU: it otherwise probes a
    # numeric CUDA device before ZeroGPU has allocated a worker.
    sys.argv = [sys.argv[0], "--cpu", "--disable-dynamic-vram",
                "--use-pytorch-cross-attention"]
    try:
        import comfy.model_management as mm
        import comfy.sd as sd
        import comfy.utils
        import nodes
        from comfy_extras.nodes_qwen import TextEncodeQwenImage21
    finally:
        sys.argv = saved_argv

    # Use a device without an index: ZeroGPU assigns it per request.
    cuda = torch.device("cuda")
    mm.args.cpu = False
    mm.args.gpu_only = True
    mm.cpu_state = mm.CPUState.GPU
    mm.vram_state = mm.VRAMState.HIGH_VRAM
    mm.get_torch_device = lambda: torch.device("cuda")

    def download(repo: str, filename: str, revision: str) -> str:
        return hf_hub_download(repo, filename, revision=revision)

    print("Loading Noct Q V4 INT8 convrot diffusion model", flush=True)
    model = sd.load_diffusion_model(
        download(MODEL_ID, CHECKPOINT, MODEL_REVISION),
        model_options={"dtype": torch.bfloat16, "load_device": cuda,
                       "offload_device": torch.device("cpu")},
        disable_dynamic=True,
    )
    if model is None:
        raise RuntimeError("ComfyUI did not recognize the Noct Q checkpoint.")
    # Comfy's quantized tensor subclasses keep internal INT8 data and
    # rotation/scaling tensors. Keep those on CPU until the real GPU worker
    # moves them: ZeroGPU's eager fake-CUDA packing does not preserve them.
    model.offload_device = cuda
    gc.collect()

    print("Loading Qwen 3 VL INT8 convrot text encoder", flush=True)
    clip = sd.load_clip(
        [download(COMPANION_ID,
                  "text_encoders/qwen3vl_8b_int8_convrot.safetensors",
                  COMPANION_REVISION)],
        clip_type=sd.CLIPType.QWEN_IMAGE,
        model_options={"dtype": torch.bfloat16, "load_device": cuda,
                       "offload_device": torch.device("cpu"),
                       "initial_device": torch.device("cpu")},
        disable_dynamic=True,
    )
    clip.patcher.offload_device = cuda
    gc.collect()

    print("Loading Qwen Image 2.1 BF16 VAE", flush=True)
    weights, metadata = comfy.utils.load_torch_file(
        download(COMPANION_ID, "vae/qwen_image_2.1_vae_bf16.safetensors",
                 COMPANION_REVISION), return_metadata=True)
    vae = sd.VAE(sd=weights, device=cuda, dtype=torch.bfloat16,
                 metadata=metadata)
    del weights
    # Keep all three components together on the ordinary Comfy loading path.
    vae.patcher.offload_device = cuda
    gc.collect()
    print("Noct Q V4, text encoder and VAE loaded; ready for GPU requests.",
          flush=True)
    return model, clip, vae, nodes, TextEncodeQwenImage21