from __future__ import annotations

import hashlib
import importlib.util
import json
import sys
from pathlib import Path
from urllib.request import urlopen

import torch
import torch.nn.functional as functional
from accelerate import init_empty_weights
from accelerate.utils import set_module_tensor_to_device
from huggingface_hub import hf_hub_download
from safetensors import safe_open
from torch import nn


BASE_MODEL = "krea/Krea-2-Turbo"
INT8_REPO = "Comfy-Org/Krea-2"
INT8_FILE = "diffusion_models/krea2_turbo_int8_convrot.safetensors"
CONVROT_COMMIT = "48a88b2fde88e986c6444fa1f51589b6089d04f3"
CONVROT_FILES = {
    "convrot.py": "89eff64518dead6ef49f894afefcf70241ce0e6eea2dfbf435f6abd9d672f6d1",
    "int8_fused_kernel.py": "cb101f0cc3d4052b791e9f579096147a4d46e334d30e174c25a0cbf9afddcba0",
}
CHECKPOINT_PREFIX = "model.diffusion_model."
SM120_MATMUL_PROFILE = {
    "BLOCK_M": 128,
    "BLOCK_N": 256,
    "BLOCK_K": 64,
    "GROUP_SIZE_M": 8,
    "num_warps": 8,
    "num_stages": 3,
}


def _download_convrot_runtime() -> Path:
    root = Path.home() / ".cache" / "anypaint" / "convrot" / CONVROT_COMMIT
    root.mkdir(parents=True, exist_ok=True)
    for name, expected_hash in CONVROT_FILES.items():
        path = root / name
        if not path.exists() or hashlib.sha256(path.read_bytes()).hexdigest() != expected_hash:
            url = (
                "https://raw.githubusercontent.com/BobJohnson24/ComfyUI-INT8-Fast/"
                f"{CONVROT_COMMIT}/{name}"
            )
            payload = urlopen(url, timeout=60).read()
            actual_hash = hashlib.sha256(payload).hexdigest()
            if actual_hash != expected_hash:
                raise RuntimeError(f"Unexpected SHA256 for {url}: {actual_hash}")
            path.write_bytes(payload)
    return root


def _load_pipeline_module(repo_id: str):
    path = Path(hf_hub_download(repo_id, "pipeline.py"))
    spec = importlib.util.spec_from_file_location("anypaint_krea2_pipeline", path)
    if spec is None or spec.loader is None:
        raise ImportError(f"Cannot import Krea 2 pipeline from {path}")
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


def _original_to_diffusers_key(key: str) -> str:
    direct = {
        "first.bias": "img_in.bias",
        "first.weight": "img_in.weight",
        "last.linear.bias": "final_layer.linear.bias",
        "last.linear.weight": "final_layer.linear.weight",
        "last.modulation.lin": "final_layer.scale_shift_table",
        "last.norm.scale": "final_layer.norm.weight",
        "tmlp.0.bias": "time_embed.linear_1.bias",
        "tmlp.0.weight": "time_embed.linear_1.weight",
        "tmlp.2.bias": "time_embed.linear_2.bias",
        "tmlp.2.weight": "time_embed.linear_2.weight",
        "tproj.1.bias": "time_mod_proj.bias",
        "tproj.1.weight": "time_mod_proj.weight",
        "txtfusion.projector.weight": "text_fusion.projector.weight",
        "txtmlp.0.scale": "txt_in.norm.weight",
        "txtmlp.1.bias": "txt_in.linear_1.bias",
        "txtmlp.1.weight": "txt_in.linear_1.weight",
        "txtmlp.3.bias": "txt_in.linear_2.bias",
        "txtmlp.3.weight": "txt_in.linear_2.weight",
    }
    if key in direct:
        return direct[key]

    attention = {"gate": "to_gate", "wk": "to_k", "wo": "to_out.0", "wq": "to_q", "wv": "to_v"}
    feed_forward = {"down": "down", "gate": "gate", "up": "up"}
    parts = key.split(".")
    if parts[0] == "blocks" and len(parts) >= 4:
        prefix, suffix = f"transformer_blocks.{parts[1]}", parts[2:]
    elif parts[0] == "txtfusion" and parts[1] in {"layerwise_blocks", "refiner_blocks"}:
        prefix, suffix = f"text_fusion.{parts[1]}.{parts[2]}", parts[3:]
    else:
        raise KeyError(f"Unsupported original Krea 2 key: {key}")
    if suffix[:2] == ["attn", "qknorm"]:
        norm_name = {"knorm": "norm_k", "qnorm": "norm_q"}.get(suffix[2])
        if norm_name and suffix[3:] == ["scale"]:
            return f"{prefix}.attn.{norm_name}.weight"
    if suffix[0] == "attn" and suffix[1] in attention and suffix[2:] == ["weight"]:
        return f"{prefix}.attn.{attention[suffix[1]]}.weight"
    if suffix[0] == "mlp" and suffix[1] in feed_forward and suffix[2:] == ["weight"]:
        return f"{prefix}.ff.{feed_forward[suffix[1]]}.weight"
    if suffix == ["mod", "lin"]:
        return f"{prefix}.scale_shift_table"
    if suffix == ["prenorm", "scale"]:
        return f"{prefix}.norm1.weight"
    if suffix == ["postnorm", "scale"]:
        return f"{prefix}.norm2.weight"
    raise KeyError(f"Unsupported original Krea 2 key: {key}")


def _pin_sm120_profile(kernel) -> None:
    expected = {key: SM120_MATMUL_PROFILE[key] for key in ("BLOCK_M", "BLOCK_N", "BLOCK_K", "GROUP_SIZE_M")}
    matches = [
        config
        for config in kernel.configs
        if config.kwargs == expected
        and config.num_warps == SM120_MATMUL_PROFILE["num_warps"]
        and config.num_stages == SM120_MATMUL_PROFILE["num_stages"]
    ]
    if len(matches) != 1:
        raise RuntimeError("Pinned ConvRot runtime lacks the validated SM120 profile")
    kernel.configs = matches
    kernel.keys = []
    kernel.cache.clear()


def _load_int8_transformer(pipeline_module):
    runtime_root = _download_convrot_runtime()
    sys.path.insert(0, str(runtime_root))
    from convrot import build_hadamard, rotate_activation
    import int8_fused_kernel

    _pin_sm120_profile(int8_fused_kernel._int8_matmul_dequant_per_row_kernel)
    linear_kernel = int8_fused_kernel.triton_int8_linear_per_row

    class Int8ConvRotLinear(nn.Module):
        def __init__(self, source: nn.Linear, group_size: int):
            super().__init__()
            self.in_features = source.in_features
            self.out_features = source.out_features
            self.group_size = group_size
            self.weight = nn.Parameter(
                torch.empty(source.out_features, source.in_features, dtype=torch.int8, device="meta"),
                requires_grad=False,
            )
            self.register_buffer("weight_scale", torch.empty(source.out_features, 1, dtype=torch.float32, device="meta"))
            self.bias = None

        def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
            shape = hidden_states.shape
            hidden_states = hidden_states.reshape(-1, shape[-1]).to(torch.bfloat16)
            hadamard = build_hadamard(self.group_size, device=hidden_states.device, dtype=hidden_states.dtype)
            hidden_states = rotate_activation(hidden_states, hadamard, self.group_size)
            output = linear_kernel(
                hidden_states,
                self.weight,
                self.weight_scale,
                self.bias,
                torch.bfloat16,
            )
            return output.reshape(*shape[:-1], self.out_features)

    config_path = hf_hub_download(BASE_MODEL, "transformer/config.json")
    config = json.loads(Path(config_path).read_text(encoding="utf-8"))
    config = {key: value for key, value in config.items() if not key.startswith("_")}
    with init_empty_weights():
        model = pipeline_module.Krea2Transformer2DModel(**config)

    checkpoint = hf_hub_download(INT8_REPO, INT8_FILE)

    def replace_module(name: str, replacement: nn.Module) -> None:
        parent_name, _, child_name = name.rpartition(".")
        parent = model.get_submodule(parent_name) if parent_name else model
        setattr(parent, child_name, replacement)

    with safe_open(checkpoint, framework="pt", device="cpu") as weights:
        raw_keys = set(weights.keys())
        tensor_map: dict[str, str] = {}
        quantized_modules: list[str] = []
        for raw_key in sorted(raw_keys):
            source_key = raw_key.removeprefix(CHECKPOINT_PREFIX)
            if source_key.endswith(".comfy_quant"):
                continue
            if source_key.endswith(".weight_scale"):
                source_weight = source_key.removesuffix(".weight_scale") + ".weight"
                target_weight = _original_to_diffusers_key(source_weight)
                target_key = target_weight.removesuffix(".weight") + ".weight_scale"
            else:
                target_key = _original_to_diffusers_key(source_key)
            tensor_map[target_key] = raw_key

        for raw_key in sorted(raw_keys):
            source_key = raw_key.removeprefix(CHECKPOINT_PREFIX)
            if not source_key.endswith(".weight") or str(weights.get_slice(raw_key).get_dtype()) != "I8":
                continue
            quant_key = raw_key.removesuffix(".weight") + ".comfy_quant"
            quant_config = json.loads(bytes(weights.get_tensor(quant_key).tolist()).decode("utf-8"))
            target_weight = _original_to_diffusers_key(source_key)
            module_name = target_weight.removesuffix(".weight")
            source_module = model.get_submodule(module_name)
            replace_module(module_name, Int8ConvRotLinear(source_module, int(quant_config["convrot_groupsize"])))
            quantized_modules.append(module_name)

        model_state = model.state_dict()
        if set(model_state) != set(tensor_map):
            raise RuntimeError("INT8 ConvRot checkpoint does not match the Krea 2 transformer layout")
        for target_key, expected in model_state.items():
            value = weights.get_tensor(tensor_map[target_key])
            if value.shape != expected.shape:
                value = value.reshape(expected.shape)
            dtype = None if not torch.is_floating_point(value) or target_key.endswith(".weight_scale") else torch.bfloat16
            set_module_tensor_to_device(model, target_key, "cuda", value=value, dtype=dtype)

    model.eval().requires_grad_(False)
    if len(quantized_modules) != 224:
        raise RuntimeError(f"Expected 224 ConvRot linear layers, got {len(quantized_modules)}")
    return model


def install_fixed_lora(transformer, repo_id: str, weight_name: str, scale: float = 1.0):
    checkpoint = hf_hub_download(repo_id, weight_name)
    pairs = []
    with safe_open(checkpoint, framework="pt", device="cpu") as weights:
        keys = set(weights.keys())
        modules = sorted(
            key.removeprefix("diffusion_model.").removesuffix(".lora_A.weight")
            for key in keys
            if key.startswith("diffusion_model.") and key.endswith(".lora_A.weight")
        )
        for source_module in modules:
            a_key = f"diffusion_model.{source_module}.lora_A.weight"
            b_key = f"diffusion_model.{source_module}.lora_B.weight"
            if b_key not in keys:
                raise RuntimeError(f"Missing LoRA B tensor for {source_module}")
            target_module = _original_to_diffusers_key(f"{source_module}.weight").removesuffix(".weight")
            a = weights.get_tensor(a_key).to("cuda", dtype=torch.bfloat16)
            b = weights.get_tensor(b_key).to("cuda", dtype=torch.bfloat16)
            pairs.append((target_module, a, b))

    hooks = []
    for module_name, a, b in pairs:
        module = transformer.get_submodule(module_name)

        def add_residual(_module, args, output, *, lora_a=a, lora_b=b):
            hidden_states = args[0].to(dtype=torch.bfloat16)
            residual = functional.linear(functional.linear(hidden_states, lora_a), lora_b)
            return output + residual.to(dtype=output.dtype) * float(scale)

        hooks.append(module.register_forward_hook(add_residual))
    if len(hooks) != 256:
        raise RuntimeError(f"Expected 256 AnyPaint LoRA modules, got {len(hooks)}")
    return hooks


def build_quantized_pipeline(repo_id: str):
    from diffusers import DiffusionPipeline

    pipeline_module = _load_pipeline_module(repo_id)
    transformer = _load_int8_transformer(pipeline_module)
    pipe = DiffusionPipeline.from_pretrained(
        BASE_MODEL,
        custom_pipeline=repo_id,
        trust_remote_code=True,
        transformer=transformer,
        torch_dtype=torch.bfloat16,
    )
    pipe.text_encoder.to("cuda")
    if torch.cuda.get_device_properties(0).total_memory < 40 * 2**30:
        pipe.vae.enable_tiling()
        pipe.vae.enable_slicing()
    pipe.vae.to("cuda")
    pipe._anypaint_runtime = {
        "dit": "INT8 ConvRot W8A8",
        "qwen": "BF16",
        "vae": "BF16",
        "convrot_commit": CONVROT_COMMIT,
        "sm120_profile": SM120_MATMUL_PROFILE,
    }
    return pipe