"""Official distilled checkpoint; FP8 transformer storage with BF16 compute.
Load one transformer shard at a time to avoid retaining 32.8 GB on disk.
"""
import gc
import json
import logging
import tempfile
from pathlib import Path
import torch
from accelerate import init_empty_weights
from accelerate.utils import set_module_tensor_to_device
from diffusers import ModularPipeline, WanAnimate2Transformer3DModel
from huggingface_hub import hf_hub_download
from safetensors import safe_open
from attention_compat import install_zero_gpu_attention

LOG = logging.getLogger(__name__)
MODEL_ID = "Wan-AI/Wan2.2-Animate-2-14B-Distilled-Diffusers"
MODEL_REVISION = "59e4141466bcb1bf9733eca1bc78be6891c9fbdf"
SKIP_FP8 = ("norm", "patch_embed", "time_embedding", "time_projection",
            "text_embedding", "img_emb", "^head")


def set_compute_dtypes(model):
    fp32_names = set(model._keep_in_fp32_modules or [])
    for name, parameter in model.named_parameters():
        dtype = torch.float32 if fp32_names.intersection(name.split(".")) else torch.bfloat16
        parameter.data = parameter.data.to(dtype=dtype)
    return model


def enable_fp8_storage(model):
    model.requires_grad_(False).eval()
    model.enable_layerwise_casting(storage_dtype=torch.float8_e4m3fn,
                                  compute_dtype=torch.bfloat16,
                                  skip_modules_pattern=SKIP_FP8)
    return model


def load_shard_into_model(model, path, expected_keys, loaded_keys, device):
    tensors = dict(model.named_parameters()) | dict(model.named_buffers())
    with safe_open(str(path), framework="pt", device="cpu") as shard:
        for name in shard.keys():
            if name not in expected_keys or name not in tensors:
                raise RuntimeError(f"Unexpected checkpoint tensor: {name}")
            if name in loaded_keys:
                raise RuntimeError(f"Duplicate checkpoint tensor: {name}")
            original, target = shard.get_tensor(name), tensors[name]
            if original.shape != target.shape:
                raise RuntimeError(f"Shape mismatch: {name}")
            # Detach even BF16 tensors from the soon-to-be-deleted mmap.
            value = original.to(dtype=target.dtype).clone()
            set_module_tensor_to_device(model, name, device, value=value, dtype=target.dtype)
            loaded_keys.add(name)
            del value, original


def load_transformer(device="cuda"):
    config = WanAnimate2Transformer3DModel.load_config(
        MODEL_ID, subfolder="transformer", revision=MODEL_REVISION)
    index_path = hf_hub_download(MODEL_ID,
        "transformer/diffusion_pytorch_model.safetensors.index.json", revision=MODEL_REVISION)
    weight_map = json.loads(Path(index_path).read_text())["weight_map"]
    with init_empty_weights():
        model = set_compute_dtypes(WanAnimate2Transformer3DModel.from_config(config))
        enable_fp8_storage(model)
    expected = set(model.state_dict())
    if set(weight_map) != expected:
        raise RuntimeError("Checkpoint does not match the pinned Diffusers model class")
    loaded = set()
    shards = sorted(set(weight_map.values()))
    for number, filename in enumerate(shards, 1):
        LOG.info("Transformer shard %d/%d: %s", number, len(shards), filename)
        with tempfile.TemporaryDirectory(prefix="wan-animate2-shard-") as folder:
            local = hf_hub_download(MODEL_ID, f"transformer/{filename}",
                                    revision=MODEL_REVISION, local_dir=folder)
            load_shard_into_model(model, local, expected, loaded, device)
        gc.collect()
    if loaded != expected or any(p.is_meta for p in model.parameters()):
        raise RuntimeError("Transformer checkpoint loading is incomplete")
    return install_zero_gpu_attention(model)


def build_pipeline():
    LOG.info("Loading Animate-2 Distilled / FP8 storage / BF16 compute")
    transformer = load_transformer()
    pipe = ModularPipeline.from_pretrained(MODEL_ID, revision=MODEL_REVISION)
    pipe.update_components(transformer=transformer)
    pipe.load_components(
        names=["text_encoder", "tokenizer", "image_encoder", "vae", "scheduler"],
        revision=MODEL_REVISION, torch_dtype=torch.bfloat16, low_cpu_mem_usage=True,
    )
    for name in ("text_encoder", "tokenizer", "image_encoder", "vae", "scheduler"):
        if getattr(pipe, name, None) is None:
            raise RuntimeError(f"Component failed to load: {name}. See startup log.")
    for name in ("text_encoder", "image_encoder", "vae"):
        getattr(pipe, name).requires_grad_(False).eval().to("cuda")
    pipe.vae.enable_tiling()
    if pipe.guider.config.guidance_scale != 1.0:
        raise RuntimeError("Expected distilled guidance_scale=1.0")
    LOG.info("READY: %s", type(pipe).__name__)
    return pipe