import json
import logging
import weakref
from pathlib import Path
from types import MethodType

import accelerate
import comfy.sd
import comfy.supported_models_base
import folder_paths
import safetensors
import torch
from comfy.ldm.lightricks.model import LTXFrequenciesPrecision, LTXRopeType
from comfy.utils import load_torch_file
from transformers import (
    AutoImageProcessor,
    Gemma3ForConditionalGeneration,
    Gemma3Processor,
)

from .embeddings_connector import Embeddings1DConnector
from .gemma_encoder import (
    GemmaFeaturesExtractorProjLinear,
    LTXVGemmaTextEncoderModel,
    find_matching_dir,
    ltxv_gemma_tokenizer,
)
from .nodes_registry import comfy_node

logger = logging.getLogger(__name__)

PREFIX_BASE = "model.diffusion_model."


def _is_device_like(x):
    return isinstance(x, (str, torch.device, int)) or isinstance(
        x, torch.Tensor
    )  # to(tensor) → matches tensor's device/dtype


def _is_dtype_like(x):
    # torch.float16 / torch.bfloat16 / torch.float32, etc.
    return isinstance(x, torch.dtype)


def protect_module_transfers(module: torch.nn.Module, warn=False):
    """
    Make .to/.cuda/.cpu NO-OPs for device moves across an entire module tree,
    while still allowing dtype/memory_format conversions.
    """

    def wrap_one(m: torch.nn.Module):
        if hasattr(m, "__to_patched"):  # idempotent
            return

        orig_to = m.__class__.to
        proxy = weakref.proxy(m)

        def safe_to(self, *args, **kwargs):
            # Block explicit/implicit device moves
            device_in_kwargs = "device" in kwargs
            device_like_positional = len(args) > 0 and _is_device_like(args[0])

            # Allow pure dtype/memory_format conversions:
            #   .to(torch.float16)
            #   .to(dtype=torch.float16)
            #   .to(memory_format=...)

            if device_in_kwargs or device_like_positional:
                if warn:
                    print("[protect_module_transfers] Ignoring device move .to(...).")
                return self  # NO-OP for device moves

            # Otherwise (dtype / memory_format / non_blocking without device) → delegate
            return orig_to(self, *args, **kwargs)

        # Bind methods
        m.to = MethodType(safe_to, proxy)
        m.cuda = MethodType(lambda self, *a, **k: self, proxy)  # NO-OP
        m.cpu = MethodType(lambda self, *a, **k: self, proxy)  # NO-OP
        m.__to_patched = True

    # Recurse over the whole tree
    for sub in module.modules():  # includes module itself and children
        wrap_one(sub)

    return module


def load_proj_matrix_from_ltxv(ltxv_path: Path, prefix=""):
    with safetensors.safe_open(ltxv_path, framework="pt", device="cpu") as f:
        keys = filter(lambda key: key.startswith(prefix), f.keys())
        sd = {k.removeprefix(prefix): f.get_tensor(k) for k in keys if k in f.keys()}

    if not sd:
        return None

    with torch.device("meta"):
        model = GemmaFeaturesExtractorProjLinear()

    for name, tensor in sd.items():
        accelerate.utils.set_module_tensor_to_device(model, name, "cpu", value=tensor)

    device_map = accelerate.infer_auto_device_map(
        model,
        max_memory=None,
    )

    model = accelerate.dispatch_model(
        model,
        device_map=device_map,
    )

    return model


def load_proj_matrix_from_checkpoint(checkpoint_path: Path):
    """
    Load model weights from a checkpoint file.
    :param checkpoint_path: Path to the checkpoint file.
    """
    with torch.device("meta"):
        model = GemmaFeaturesExtractorProjLinear()
    model = accelerate.load_checkpoint_and_dispatch(
        model, str(checkpoint_path.resolve()), device_map="auto"
    )
    return model


def ltxv_gemma_clip(encoder_path, ltxv_path, processor=None, dtype=None):
    class _LTXVGemmaTextEncoderModel(LTXVGemmaTextEncoderModel):
        def __init__(self, device="cpu", dtype=dtype, model_options={}):
            dtype = torch.bfloat16  # TODO: make this configurable
            gemma_model = Gemma3ForConditionalGeneration.from_pretrained(
                encoder_path,
                local_files_only=True,
                torch_dtype=dtype,
                device_map="auto",
            )

            feature_extractor_linear = load_proj_matrix_from_ltxv(
                ltxv_path, "text_embedding_projection."
            )
            if feature_extractor_linear is None:
                feature_extractor_linear = load_proj_matrix_from_checkpoint(
                    encoder_path / "proj_linear.safetensors"
                )

            embeddings_connector = load_video_embeddings_connector(ltxv_path)
            audio_embeddings_connector = load_audio_embeddings_connector(ltxv_path)
            super().__init__(
                model=gemma_model,
                feature_extractor_linear=feature_extractor_linear,
                embeddings_connector=embeddings_connector,
                audio_embeddings_connector=audio_embeddings_connector,
                processor=processor,
                dtype=dtype,
                device=device,
            )

            # Due to memory management issue on mgpu we need to disable the offloading of sharded gemma model.
            protect_module_transfers(self)

    return _LTXVGemmaTextEncoderModel


def load_video_embeddings_connector(ltxv_path, dtype=torch.bfloat16):
    sd, metadata = load_torch_file(str(ltxv_path), return_metadata=True)
    config = json.loads(metadata.get("config", "{}"))
    transformer_config = config.get("transformer", {})
    rope_type = LTXRopeType.from_dict(transformer_config)
    frequencies_precision = LTXFrequenciesPrecision.from_dict(transformer_config)
    pe_max_pos = transformer_config.get("connector_positional_embedding_max_pos", [1])
    sd_keys = list(sd.keys())

    video_only_connector_prefix = f"{PREFIX_BASE}embeddings_connector."
    av_connector_prefix = f"{PREFIX_BASE}video_embeddings_connector."
    prefix = (
        av_connector_prefix
        if f"{PREFIX_BASE}audio_adaln_single.linear.weight" in sd_keys
        else video_only_connector_prefix
    )
    return load_embeddings_connector(
        sd, prefix, dtype, rope_type, frequencies_precision, pe_max_pos
    )


def load_audio_embeddings_connector(ltxv_path, dtype=torch.bfloat16):
    sd, metadata = load_torch_file(str(ltxv_path), return_metadata=True)
    config = json.loads(metadata.get("config", "{}"))
    transformer_config = config.get("transformer", {})
    rope_type = LTXRopeType.from_dict(transformer_config)
    frequencies_precision = LTXFrequenciesPrecision.from_dict(transformer_config)
    pe_max_pos = transformer_config.get("connector_positional_embedding_max_pos", [1])
    return load_embeddings_connector(
        sd,
        f"{PREFIX_BASE}audio_embeddings_connector.",
        dtype,
        rope_type,
        frequencies_precision,
        pe_max_pos,
    )


def load_embeddings_connector(
    sd,
    connector_prefix,
    dtype=torch.bfloat16,
    rope_type=LTXRopeType.INTERLEAVED,
    frequencies_precision=LTXFrequenciesPrecision.FLOAT32,
    pe_max_pos=None,
):
    sd_connector = {
        k[len(connector_prefix) :]: v
        for k, v in sd.items()
        if k.startswith(connector_prefix)
    }

    if len(sd_connector) == 0:
        return None

    operations = comfy.ops.pick_operations(dtype, dtype, disable_fast_fp8=True)
    connector = Embeddings1DConnector(
        dtype=dtype,
        operations=operations,
        positional_embedding_max_pos=pe_max_pos if pe_max_pos is not None else [1],
        split_rope=rope_type == LTXRopeType.SPLIT,
        double_precision_rope=frequencies_precision == LTXFrequenciesPrecision.FLOAT64,
    )
    connector.load_state_dict(sd_connector)
    device_map = accelerate.infer_auto_device_map(connector)
    connector = accelerate.dispatch_model(connector, device_map=device_map)
    return connector


@comfy_node(
    name="LTXVGemmaCLIPModelLoaderMGPU", description="Gemma 3 Model Loader on MGPU"
)
class LTXVGemmaCLIPModelLoaderMGPU:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "gemma_path": (
                    folder_paths.get_filename_list("text_encoders"),
                    {"tooltip": "The name of the text encoder model to load."},
                ),
                "ltxv_path": (
                    folder_paths.get_filename_list("checkpoints"),
                    {"tooltip": "The name of the ltxv model to load."},
                ),
                "max_length": (
                    "INT",
                    {"default": 1024, "min": 16, "max": 131072, "step": 8},
                ),
            }
        }

    RETURN_TYPES = ("CLIP",)
    RETURN_NAMES = ("clip",)
    FUNCTION = "load_model"
    CATEGORY = "lightricks/LTXV"
    TITLE = "LTXV Gemma CLIP Loader"
    OUTPUT_NODE = False

    def load_model(self, gemma_path: str, ltxv_path: str, max_length: int):
        path = Path(folder_paths.get_full_path("text_encoders", gemma_path))
        model_root = path.parents[1]
        tokenizer_path = Path(find_matching_dir(model_root, "tokenizer.model"))
        gemma_model_path = Path(find_matching_dir(model_root, "model*.safetensors"))

        tokenizer_class = ltxv_gemma_tokenizer(tokenizer_path, max_length=max_length)

        processor = None
        try:
            image_processor = AutoImageProcessor.from_pretrained(
                str(model_root),
                local_files_only=True,
            )
            processor = Gemma3Processor(
                image_processor=image_processor,
                tokenizer=tokenizer_class().tokenizer,
            )
            logger.info(f"Loaded processor from {model_root} - enhancement enabled")
        except Exception as e:
            logger.warning(f"Could not load processor from {model_root}: {e}")

        clip_dtype = torch.bfloat16
        ltxv_full_path = folder_paths.get_full_path("checkpoints", ltxv_path)
        clip_target = comfy.supported_models_base.ClipTarget(
            tokenizer=tokenizer_class,
            clip=ltxv_gemma_clip(
                gemma_model_path, ltxv_full_path, processor=processor, dtype=clip_dtype
            ),
        )

        return (comfy.sd.CLIP(clip_target),)