"""Load the tiny LTX-2.5 split checkpoint into a Diffusers pipeline."""

import json
import tempfile
from pathlib import Path

import torch
from diffusers import (
    AutoencoderKLLTX2Audio,
    AutoencoderKLLTX2Video,
    FlowMatchEulerDiscreteScheduler,
    LTX2Pipeline,
    LTX2VideoTransformer3DModel,
)
from diffusers.loaders.single_file_utils import (
    convert_ltx2_audio_vae_to_diffusers,
    convert_ltx2_transformer_to_diffusers,
    convert_ltx2_vae_to_diffusers,
)
from diffusers.pipelines.ltx2.connectors import LTX2TextConnectors
from diffusers.pipelines.ltx2.vocoder import LTX2Vocoder
from safetensors.torch import load_file
from transformers import (
    AutoTokenizer,
    Gemma4UnifiedConfig,
    Gemma4UnifiedForConditionalGeneration,
    Gemma4UnifiedTextConfig,
    Gemma4UnifiedVisionConfig,
)

FILES = {
    "transformer": "diffusion_models/ltx-2.5-22b-distilled-transformer-bf16.safetensors",
    "text_encoder": "text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors",
    "vae": "vae/ltx-2.5-video-vae-bf16.safetensors",
    "audio_vae": "vae/ltx-2.5-audio-vae-bf16.safetensors",
}
ASSETS = {
    "tokenizer_json": "tokenizer.json",
    "hf_asset__tokenizer_config.json": "tokenizer_config.json",
    "hf_asset__chat_template.jinja": "chat_template.jinja",
}
# Diffusers name -> Comfy single-file name for the tensors the converters drop.
CONNECTOR_PREFIX = {
    "video_connector.transformer_blocks.": "video_embeddings_connector.transformer_1d_blocks.",
    "audio_connector.transformer_blocks.": "audio_embeddings_connector.transformer_1d_blocks.",
    "video_connector.learnable_registers": "video_embeddings_connector.learnable_registers",
    "audio_connector.learnable_registers": "audio_embeddings_connector.learnable_registers",
}
TEXT_PREFIX = {
    "model.language_model.": "model.",
    "model.embed_vision.": "vision_model.",
}
PROJECTIONS = {
    "video_text_proj_in.weight": "text_embedding_projection.video_aggregate_embed.weight",
    "audio_text_proj_in.weight": "text_embedding_projection.audio_aggregate_embed.weight",
}


def _connector_name(name):
    for local, source in CONNECTOR_PREFIX.items():
        if name.startswith(local):
            name = source + name[len(local) :]
            break
    return name.replace("norm_q", "q_norm").replace("norm_k", "k_norm")


def _text_name(name):
    for local, source in TEXT_PREFIX.items():
        if name.startswith(local):
            return source + name[len(local) :]
    return None


def load_transformer(config, state):
    model = LTX2VideoTransformer3DModel(**config["transformer"])
    plain = {name: value for name, value in state.items() if "embeddings_connector" not in name}
    model.load_state_dict(convert_ltx2_transformer_to_diffusers(plain))
    return model


def load_connectors(config, transformer_state, text_state):
    model = LTX2TextConnectors(**config["connectors"])
    weights = {}
    for name in model.state_dict():
        if name in PROJECTIONS:
            weights[name] = text_state[PROJECTIONS[name]]
        else:
            weights[name] = transformer_state["model.diffusion_model." + _connector_name(name)]
    model.load_state_dict(weights, strict=True)
    return model


def load_text_encoder(config, text_state):
    text = Gemma4UnifiedTextConfig(
        **{key: value for key, value in config["text_encoder"].items() if key != "tie_word_embeddings"}
    )
    gemma = Gemma4UnifiedConfig(
        text_config=text.to_dict(),
        vision_config=Gemma4UnifiedVisionConfig(
            patch_size=16, pooling_kernel_size=2, mm_embed_dim=8, mm_posemb_size=16, output_proj_dims=8
        ).to_dict(),
    )
    gemma.tie_word_embeddings = True
    model = Gemma4UnifiedForConditionalGeneration(gemma)
    weights = {}
    for name in model.state_dict():
        source = _text_name(name)
        if source is not None and source in text_state:
            weights[name] = text_state[source]
    model.load_state_dict(weights, strict=False)
    return model


def load_tokenizer(text_state):
    with tempfile.TemporaryDirectory() as directory:
        for tensor_name, file_name in ASSETS.items():
            Path(directory, file_name).write_bytes(text_state[tensor_name].numpy().tobytes())
        return AutoTokenizer.from_pretrained(directory)


def load_pipeline(model_id, device=None):
    model_id = Path(model_id)
    config = json.loads((model_id / "config.json").read_text())
    transformer_state = load_file(str(model_id / FILES["transformer"]))
    text_state = load_file(str(model_id / FILES["text_encoder"]))

    transformer = load_transformer(config, transformer_state)
    connectors = load_connectors(config, transformer_state, text_state)
    vae = AutoencoderKLLTX2Video(**config["vae"])
    vae.load_state_dict(convert_ltx2_vae_to_diffusers(load_file(str(model_id / FILES["vae"]))))

    audio_state = load_file(str(model_id / FILES["audio_vae"]))
    audio_vae = AutoencoderKLLTX2Audio(**config["audio_vae"])
    audio_vae.load_state_dict(
        convert_ltx2_audio_vae_to_diffusers(
            {name: value for name, value in audio_state.items() if not name.startswith("vocoder.")}
        )
    )
    vocoder = LTX2Vocoder(**config["vocoder"])
    vocoder.load_state_dict(
        {
            name.removeprefix("vocoder."): value
            for name, value in audio_state.items()
            if name.startswith("vocoder.")
        }
    )

    text_encoder = load_text_encoder(config, text_state)
    pipeline = LTX2Pipeline(
        scheduler=FlowMatchEulerDiscreteScheduler(),
        vae=vae,
        audio_vae=audio_vae,
        text_encoder=text_encoder,
        tokenizer=load_tokenizer(text_state),
        connectors=connectors,
        transformer=transformer,
        vocoder=vocoder,
    )
    for module in (transformer, vae, audio_vae, text_encoder, connectors, vocoder):
        module.eval()
    pipeline.to(device or ("cuda" if torch.cuda.is_available() else "cpu"))
    return pipeline