"""Pinned ComfyUI-INT8-Fast integration; no GPU work at module import."""

import importlib.util
import json
from pathlib import Path
import struct
import subprocess
import sys

INT8_REPO = "https://github.com/BobJohnson24/ComfyUI-INT8-Fast.git"
# Matches the user's working OTUNetLoaderW8A8 workflow and local installation.
INT8_COMMIT = "48a88b2fde88e986c6444fa1f51589b6089d04f3"
MODEL_ALIASES = {
    "anima-turbo-v1.1.safetensors": "anima-turbo-v1.1_INT8.safetensors",
    "anima-aesthetic-v1.1.safetensors": "anima-aestheticV11_INT8.safetensors",
    "anima_aestheticV11.safetensors": "anima-aestheticV11_INT8.safetensors",
    "anima-aestheticV11.safetensors": "anima-aestheticV11_INT8.safetensors",
    "anima_aestheticV11_INT8.safetensors": "anima-aestheticV11_INT8.safetensors",
    "fnMixAnimaTurbo_baseNoTurbo.safetensors": "fnMixAnimaTurbo_baseNoTurbo_INT8.safetensors",
    "fwaunAnima_v11.safetensors": "fwaunAnima_v11_INT8.safetensors",
    "tekitoMix_v32.safetensors": "tekitoMix_v32_INT8.safetensors",
}
INT8_ANIMA_MODELS = frozenset(MODEL_ALIASES.values())


def ensure_int8_node(comfy_path):
    """Fetch the pinned node; never delete or reset an existing checkout."""
    destination = Path(comfy_path) / "custom_nodes" / "ComfyUI-INT8-Fast"
    if destination.exists():
        revision = subprocess.check_output(
            ["git", "-C", str(destination), "rev-parse", "HEAD"], text=True
        ).strip()
        dirty = subprocess.check_output(
            ["git", "-C", str(destination), "status", "--porcelain", "--untracked-files=no"], text=True
        ).strip()
        if revision != INT8_COMMIT or dirty:
            raise RuntimeError("INT8-Fast checkout differs from the pinned version; existing files were preserved.")
    else:
        destination.parent.mkdir(parents=True, exist_ok=True)
        # clone into a fresh path only; no cleanup that could erase local changes.
        subprocess.run(["git", "clone", "--no-checkout", "--filter=blob:none", INT8_REPO, str(destination)], check=True)
        subprocess.run(["git", "-C", str(destination), "checkout", "--detach", INT8_COMMIT], check=True)
        revision = subprocess.check_output(
            ["git", "-C", str(destination), "rev-parse", "HEAD"], text=True
        ).strip()
        if revision != INT8_COMMIT:
            raise RuntimeError("Unexpected INT8-Fast revision")
    if not (destination / "int8_unet_loader.py").is_file():
        raise RuntimeError("Pinned INT8-Fast loader is missing")
    return destination


def make_int8_loader(comfy_path):
    """Import the node as a package so its relative imports/layout setup run."""
    destination = ensure_int8_node(comfy_path)
    name = "_space_int8_fast"
    if name not in sys.modules:
        spec = importlib.util.spec_from_file_location(
            name, destination / "__init__.py", submodule_search_locations=[str(destination)]
        )
        module = importlib.util.module_from_spec(spec)
        sys.modules[name] = module
        try:
            spec.loader.exec_module(module)
        except BaseException:
            # Do not leave a partially initialized package available on retry.
            for key in list(sys.modules):
                if key == name or key.startswith(name + "."):
                    del sys.modules[key]
            raise
    loader_class = sys.modules[name].NODE_CLASS_MAPPINGS.get("OTUNetLoaderW8A8")
    if loader_class is None:
        raise RuntimeError("INT8-Fast could not import OTUNetLoaderW8A8; check ComfyUI/PyTorch dependencies.")
    return loader_class()


def validate_int8_file(path):
    """Inspect only the safetensors header, not multi-GB weights or metadata."""
    with Path(path).open("rb") as stream:
        prefix = stream.read(8)
        if len(prefix) != 8:
            raise ValueError("INT8 model has an incomplete safetensors header")
        length = struct.unpack("<Q", prefix)[0]
        if not 2 <= length <= 16_000_000:
            raise ValueError("INT8 model has an invalid safetensors header length")
        data = stream.read(length)
        if len(data) != length:
            raise ValueError("INT8 model has a truncated safetensors header")
    header = json.loads(data)
    tensors = {key: value for key, value in header.items() if key != "__metadata__"}
    if not any(value.get("dtype") == "I8" for value in tensors.values()):
        raise ValueError("Registered INT8 model contains no INT8 weights; refusing BF16 fallback")
    if not any(key.endswith(".comfy_quant") for key in tensors):
        raise ValueError("Registered INT8 model is missing INT8-Fast quantization settings")


def load_anima_model(name, standard_loader, int8_loader=None, model_directory=None):
    if name not in INT8_ANIMA_MODELS:
        return standard_loader.load_unet(unet_name=name, weight_dtype="default")
    if int8_loader is None or model_directory is None:
        raise RuntimeError("INT8-Fast loader is required for the registered INT8 model")
    validate_int8_file(Path(model_directory) / name)
    # Match image_anima_INT8_TurboSwitch.json. Never quantize at Space startup.
    return int8_loader.load_unet(
        unet_name=name, weight_dtype="default", model_type="anima",
        on_the_fly_quantization=False, enable_convrot=True, lora_mode="None",
    )