import json
import logging
from glob import glob
from pathlib import Path
from typing import List, Optional, Tuple

import comfy.model_management
import comfy.ops
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 einops import rearrange
from PIL import Image
from transformers import (
    AutoImageProcessor,
    AutoTokenizer,
    Gemma3Config,
    Gemma3ForConditionalGeneration,
    Gemma3Processor,
)
from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS
from transformers.models.auto.modeling_auto import MODEL_MAPPING_NAMES

from .embeddings_connector import Embeddings1DConnector
from .nodes_registry import comfy_node

logger = logging.getLogger(__name__)

PREFIX_BASE = "model.diffusion_model."


def _load_system_prompt(filename: str) -> str:
    """Load system prompt from file at module level."""
    try:
        prompt_path = Path(__file__).parent / "system_prompts" / filename
        if prompt_path.exists():
            return prompt_path.read_text(encoding="utf-8").strip()
    except Exception as e:
        logger.warning(f"Could not load {filename}: {e}")
    return ""


DEFAULT_T2V_SYSTEM_PROMPT = _load_system_prompt("gemma_t2v_system_prompt.txt")
DEFAULT_I2V_SYSTEM_PROMPT = _load_system_prompt("gemma_i2v_system_prompt.txt")


def tensor_to_pil(tensor: torch.Tensor) -> Image.Image:
    """Convert ComfyUI image tensor to PIL Image."""
    if tensor.dim() == 4:
        tensor = tensor[0]
    numpy_image = (tensor.cpu().numpy() * 255).astype("uint8")
    return Image.fromarray(numpy_image)


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)
    return connector


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
    model = GemmaFeaturesExtractorProjLinear()
    model.load_state_dict(sd)
    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.
    """
    model = GemmaFeaturesExtractorProjLinear()
    loaded_state_dict = load_torch_file(str(checkpoint_path), return_metadata=False)
    if "aggregate_embed.weight" not in loaded_state_dict:
        raise ValueError(
            f"Checkpoint {checkpoint_path} does not contain 'aggregate_embed.weight'."
        )
    model.load_state_dict(loaded_state_dict)
    return model


class LTXVGemmaTokenizer:
    def __init__(self, tokenizer_path: str, max_length: int = 1024):
        self.tokenizer = AutoTokenizer.from_pretrained(
            tokenizer_path, local_files_only=True, model_max_length=max_length
        )
        # Gemma expects left padding for chat-style prompts; for plain text it doesn't matter much.
        self.tokenizer.padding_side = "left"
        if self.tokenizer.pad_token is None:
            self.tokenizer.pad_token = self.tokenizer.eos_token

        self.max_length = max_length

    def tokenize_with_weights(self, text: str, return_word_ids: bool = False):
        text = text.strip()
        encoded = self.tokenizer(
            text,
            padding="max_length",
            max_length=self.max_length,
            truncation=True,
            return_tensors="pt",
        )
        input_ids = encoded.input_ids
        attention_mask = encoded.attention_mask
        tuples = [
            (token_id, attn, i)
            for i, (token_id, attn) in enumerate(zip(input_ids[0], attention_mask[0]))
        ]
        out = {"gemma": tuples}

        if not return_word_ids:
            out = {k: [(t, w) for t, w, _ in v] for k, v in out.items()}

        return out


class GemmaFeaturesExtractorProjLinear(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.aggregate_embed = torch.nn.Linear(3840 * 49, 3840, bias=False)

    def forward(self, x):
        return self.aggregate_embed(x)


class LTXVGemmaTextEncoderModel(torch.nn.Module):
    def __init__(
        self,
        model: Gemma3ForConditionalGeneration,
        feature_extractor_linear: GemmaFeaturesExtractorProjLinear,
        embeddings_connector: Embeddings1DConnector,
        audio_embeddings_connector: Embeddings1DConnector | None = None,
        processor: Gemma3Processor | None = None,
        dtype=torch.bfloat16,
        device="cpu",
    ):
        super().__init__()
        self.model = model
        self.processor = processor
        self.feature_extractor_linear = feature_extractor_linear.to(dtype=dtype)
        self.embeddings_connector = embeddings_connector.to(dtype=dtype)
        self.audio_embeddings_connector = (
            audio_embeddings_connector.to(dtype=dtype)
            if audio_embeddings_connector is not None
            else None
        )
        self.dtypes = set([dtype])
        # Cache an estimate of memory required to load/keep the model on device
        # weights size + small overhead
        self._model_memory_required = (
            comfy.model_management.module_size(self.model) + 256 * 1024 * 1024
        )

    def set_clip_options(self, options):
        pass

    def reset_clip_options(self):
        pass

    @staticmethod
    def norm_and_concat_padded_batch(
        encoded_text: torch.Tensor,
        sequence_lengths: torch.Tensor,
        padding_side: str = "right",
    ) -> torch.Tensor:
        """
        Normalize a 4D tensor [B, T, D, L] per sample and per layer, using sequence_lengths to mask.
        Returns [B, T,  D * L] tensor with original padding preserved.

        Args:
            encoded_text: 4D tensor [B, T, D, L]
            sequence_lengths: 1D tensor [B] with actual sequence lengths
            padding_side: "left" or "right" to indicate which side has padding
        """
        B, T, D, L = encoded_text.shape
        device = encoded_text.device

        # Build mask: [B, T, 1, 1]
        token_indices = torch.arange(T, device=device)[None, :]  # [1, T]

        if padding_side == "right":
            # For right padding, valid tokens are from 0 to sequence_length-1
            mask = token_indices < sequence_lengths[:, None]  # [B, T]
        elif padding_side == "left":
            # For left padding, valid tokens are from (T - sequence_length) to T-1
            start_indices = T - sequence_lengths[:, None]  # [B, 1]
            mask = token_indices >= start_indices  # [B, T]
        else:
            raise ValueError(
                f"padding_side must be 'left' or 'right', got {padding_side}"
            )

        mask = rearrange(mask, "b t -> b t 1 1")

        # Compute masked mean: [B, 1, 1, L]
        masked = encoded_text.masked_fill(~mask, 0.0)
        # denom = mask.sum(dim=(1, 2), keepdim=True)  # [B, 1, 1, 1]
        denom = (sequence_lengths * D).view(B, 1, 1, 1)
        mean = masked.sum(dim=(1, 2), keepdim=True) / (denom + 1e-6)

        # Compute masked min/max: [B, 1, 1, L]
        x_min = encoded_text.masked_fill(~mask, float("inf")).amin(
            dim=(1, 2), keepdim=True
        )
        x_max = encoded_text.masked_fill(~mask, float("-inf")).amax(
            dim=(1, 2), keepdim=True
        )
        range_ = x_max - x_min

        # Normalize only the valid tokens
        normed = 8 * (encoded_text - mean) / (range_ + 1e-6)

        # concat to be [Batch, T,  D * L] - this preserves the original structure
        normed = normed.reshape(B, T, -1)  # [B, T, D * L]

        # Apply mask to preserve original padding (set padded positions to 0)
        mask_flattened = rearrange(mask, "b t 1 1 -> b t 1").expand(-1, -1, D * L)
        normed = normed.masked_fill(~mask_flattened, 0.0)

        return normed

    def forward(self, input_ids, attention_mask, padding_side="right"):
        outputs = self.model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            output_hidden_states=True,
        )
        encoded_text_features = torch.stack(outputs.hidden_states, dim=-1)
        encoded_text_features_dtype = encoded_text_features.dtype

        sequence_lengths = attention_mask.sum(dim=-1)
        normed_concated_encoded_text_features = (
            LTXVGemmaTextEncoderModel.norm_and_concat_padded_batch(
                encoded_text_features, sequence_lengths, padding_side=padding_side
            )
        )

        projected = self.feature_extractor_linear(
            normed_concated_encoded_text_features.to(encoded_text_features_dtype)
        )

        return projected

    def encode_token_weights(self, token_weight_pairs):
        token_pairs = token_weight_pairs["gemma"]
        input_ids = torch.tensor(
            [[t[0] for t in token_pairs]], device=self.model.device
        )
        attention_mask = torch.tensor(
            [[w[1] for w in token_pairs]], device=self.model.device
        )

        self.to(self.model.device)

        encoded_input = self(input_ids, attention_mask, padding_side="left")
        # convert attention mask to format embeddings connector expects.
        connector_attention_mask = (attention_mask - 1).to(encoded_input.dtype).reshape(
            (attention_mask.shape[0], 1, -1, attention_mask.shape[-1])
        ) * torch.finfo(encoded_input.dtype).max
        encoded, encoded_connector_attention_mask = self.embeddings_connector(
            encoded_input,
            connector_attention_mask,
        )

        # restore the mask values to int64
        attention_mask = (encoded_connector_attention_mask < 0.000001).to(torch.int64)

        attention_mask = attention_mask.reshape([encoded.shape[0], encoded.shape[1], 1])
        encoded = encoded * attention_mask

        if self.audio_embeddings_connector is not None:
            encoded_for_audio, _ = self.audio_embeddings_connector(
                encoded_input, connector_attention_mask
            )
            encoded = torch.concat(
                [encoded, encoded_for_audio], dim=len(encoded.shape) - 1
            )

        return encoded, None, {"attention_mask": attention_mask.squeeze(-1)}

    def load_sd(self, sd):
        return self.model.load_state_dict(sd, strict=False)

    def memory_required(self, input_shape):
        # Return a conservative estimate in bytesed(input_shape)
        return self._model_memory_required


def ltxv_gemma_tokenizer(tokenizer_path, max_length=256):
    class _LTXVGemmaTokenizer(LTXVGemmaTokenizer):
        def __init__(self, embedding_directory=None, tokenizer_data={}):
            super().__init__(tokenizer_path, max_length=max_length)

    return _LTXVGemmaTokenizer


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,
            )

            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,
            )

    return _LTXVGemmaTextEncoderModel


def find_matching_dir(root_path: str, pattern: str) -> str:
    """
    Recursively search for files matching a glob pattern and return the parent directory of the first match.
    """
    matches = [
        Path(p)
        for p in glob(f"{root_path}/**", recursive=True)
        if Path(p).match(pattern)
    ]
    if not matches:
        raise FileNotFoundError(
            f"No files matching pattern '{pattern}' found under {root_path}"
        )
    return str(matches[0].parent)


@comfy_node(name="LTXVGemmaCLIPModelLoader", description="Gemma 3 Model Loader")
class LTXVGemmaCLIPModelLoader:
    @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"))
        processor_path = Path(find_matching_dir(model_root, "preprocessor_config.json"))
        tokenizer_class = ltxv_gemma_tokenizer(tokenizer_path, max_length=max_length)

        processor = None
        try:
            image_processor = AutoImageProcessor.from_pretrained(
                str(processor_path),
                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),)


_UNICODE_REPLACEMENTS = str.maketrans(
    "\u2018\u2019\u201c\u201d\u2014\u2013\u00a0\u2032\u2212", "''\"\"-- '-"
)


def clean_response(text):
    text = text.translate(_UNICODE_REPLACEMENTS)

    # Remove leading non-letter characters
    for i, char in enumerate(text):
        if char.isalpha():
            return text[i:]
    return text


@comfy_node(name="LTXVGemmaEnhancePrompt", description="Gemma 3 Prompt Enhancer")
class LTXVGemmaEnhancePrompt:
    """Enhance prompts using Gemma 3 model. Supports T2V and I2V modes."""

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "clip": ("CLIP",),
                "prompt": ("STRING", {"multiline": True, "default": ""}),
                "system_prompt": (
                    "STRING",
                    {
                        "multiline": True,
                        "default": DEFAULT_T2V_SYSTEM_PROMPT,
                    },
                ),
                "max_tokens": (
                    "INT",
                    {"default": 512, "min": 32, "max": 1024, "step": 16},
                ),
                "bypass_i2v": ("BOOLEAN", {"default": False}),
            },
            "optional": {
                "image": ("IMAGE",),
                "seed": (
                    "INT",
                    {"default": 42, "min": 0, "max": 0xFFFFFFFF},
                ),
            },
        }

    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("enhanced_prompt",)
    FUNCTION = "enhance"
    CATEGORY = "lightricks/LTXV"
    TITLE = "LTXV Gemma Enhance Prompt"
    OUTPUT_NODE = True
    DESCRIPTION = (
        "Enhance text prompts using Gemma 3 VLLM for improved video generation."
    )

    def enhance(
        self,
        clip,
        prompt: str,
        system_prompt: str,
        max_tokens: int,
        bypass_i2v: bool,
        image: Optional[torch.Tensor] = None,
        seed: int = 42,
    ):
        if not isinstance(seed, int):
            seed = 42

        clip.load_model()
        encoder = clip.cond_stage_model

        if not hasattr(encoder, "processor") or encoder.processor is None:
            if hasattr(encoder, "gemma3_12b"):
                model, processor = transformers_gemma3_from_encoder(encoder)
            else:
                raise ValueError(
                    "Processor not loaded - enhancement not available. "
                    "Ensure your model directory has chat_template.json, processor_config.json, "
                    "and preprocessor_config.json files."
                )
        elif isinstance(encoder, LTXVGemmaTextEncoderModel):
            model = encoder.model
            processor = encoder.processor
        # Determine mode: use I2V if image is provided and not bypassed
        use_i2v = image is not None and not bypass_i2v

        # Auto-select the appropriate system prompt if user is using default T2V prompt
        if use_i2v and system_prompt.strip() == DEFAULT_T2V_SYSTEM_PROMPT.strip():
            system_prompt = DEFAULT_I2V_SYSTEM_PROMPT
            logger.info("Auto-selected I2V system prompt for image-to-video mode")

        if not system_prompt or not system_prompt.strip():
            raise ValueError(
                "system_prompt is required and cannot be empty or whitespace-only"
            )

        if use_i2v:
            pil_image = tensor_to_pil(image)
            enhanced_prompt = enhance_i2v(
                processor=processor,
                model=model,
                prompt=prompt,
                image=pil_image,
                system_prompt=system_prompt,
                max_new_tokens=max_tokens,
                seed=seed,
            )
        else:
            enhanced_prompt = enhance_t2v(
                processor=processor,
                model=model,
                prompt=prompt,
                system_prompt=system_prompt,
                max_new_tokens=max_tokens,
                seed=seed,
            )

        enhanced_prompt = clean_response(enhanced_prompt)

        return (enhanced_prompt,)


def _enhance(
    processor: Gemma3Processor,
    model: Gemma3ForConditionalGeneration,
    messages: list,
    image: Optional[Image.Image] = None,
    max_new_tokens: int = 512,
    seed: int = 42,
) -> str:
    """Common enhancement logic for both T2V and I2V modes."""
    if processor is None:
        raise ValueError("Processor not loaded - enhancement not available")

    text = processor.tokenizer.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )

    model_inputs = processor(
        text=text,
        images=image,
        return_tensors="pt",
    ).to(model.device)

    pad_token_id = (
        processor.tokenizer.pad_token_id
        if processor.tokenizer.pad_token_id is not None
        else 0
    )
    model_inputs = _pad_inputs_for_attention_alignment(model_inputs, pad_token_id)

    with (
        torch.inference_mode(),
        torch.random.fork_rng(devices=[model.device]),
        torch.autocast(device_type=model.device.type, dtype=model.dtype),
    ):
        torch.manual_seed(seed)
        outputs = model.generate(
            **model_inputs,
            max_new_tokens=max_new_tokens,
            do_sample=True,
            temperature=0.7,
        )
        generated_ids = outputs[0][len(model_inputs.input_ids[0]) :]
        enhanced_prompt = processor.tokenizer.decode(
            generated_ids, skip_special_tokens=True
        )

    return enhanced_prompt


def enhance_t2v(
    processor: Gemma3Processor,
    model: Gemma3ForConditionalGeneration,
    prompt: str,
    system_prompt: str,
    max_new_tokens: int = 512,
    seed: int = 42,
) -> str:
    """Enhance a text prompt for T2V generation."""
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": f"User Raw Input Prompt: {prompt}."},
    ]
    return _enhance(
        processor, model, messages, max_new_tokens=max_new_tokens, seed=seed
    )


def enhance_i2v(
    processor: Gemma3Processor,
    model: Gemma3ForConditionalGeneration,
    prompt: str,
    image: Image.Image,
    system_prompt: str,
    max_new_tokens: int = 512,
    seed: int = 42,
) -> str:
    """Enhance a text prompt for I2V generation using a reference image."""
    messages = [
        {"role": "system", "content": system_prompt},
        {
            "role": "user",
            "content": [
                {"type": "image"},
                {"type": "text", "text": f"User Raw Input Prompt: {prompt}."},
            ],
        },
    ]
    return _enhance(
        processor,
        model,
        messages,
        image=image,
        max_new_tokens=max_new_tokens,
        seed=seed,
    )


def _cat_with_padding(
    tensor: torch.Tensor,
    padding_length: int,
    value: int | float,
) -> torch.Tensor:
    """Concatenate a tensor with a padding tensor of the given value."""
    return torch.cat(
        [
            tensor,
            torch.full(
                (1, padding_length),
                value,
                dtype=tensor.dtype,
                device=tensor.device,
            ),
        ],
        dim=1,
    )


def _pad_inputs_for_attention_alignment(model_inputs, pad_token_id, alignment: int = 8):
    """Pad sequence length to multiple of alignment for Flash Attention compatibility.

    Flash Attention within SDPA requires sequence lengths aligned to 8 bytes.
    This pads input_ids, attention_mask, and token_type_ids (if present) to prevent
    'p.attn_bias_ptr is not correctly aligned' errors.
    """
    seq_len = model_inputs.input_ids.shape[1]
    padded_len = ((seq_len + alignment - 1) // alignment) * alignment
    padding_length = padded_len - seq_len

    if padding_length > 0:
        model_inputs["input_ids"] = _cat_with_padding(
            model_inputs.input_ids, padding_length, pad_token_id
        )

        model_inputs["attention_mask"] = _cat_with_padding(
            model_inputs.attention_mask, padding_length, 0
        )

        if (
            "token_type_ids" in model_inputs
            and model_inputs["token_type_ids"] is not None
        ):
            model_inputs["token_type_ids"] = _cat_with_padding(
                model_inputs["token_type_ids"], padding_length, 0
            )

    return model_inputs


def _locate_model_within_model(super_model, model_name):
    class_name = MODEL_MAPPING_NAMES.get(model_name, None)
    if class_name is None:
        return None
    for module in super_model.modules():
        if module.__class__.__name__ == class_name:
            return module
    return None


def _locate_unique_parameter_owner_by_leaf(
    root: torch.nn.Module,
    leaf_param_name: str,
    must_have_in_path: Optional[str] = None,
) -> Optional[Tuple[torch.nn.Module, str, torch.nn.Parameter, str]]:

    modules = dict(root.named_modules())

    candidates: List[Tuple[torch.nn.Module, str, torch.nn.Parameter, str]] = []
    for full_name, p in root.named_parameters(recurse=True):
        parts = full_name.split(".")
        leaf = parts[-1]
        if leaf != leaf_param_name:
            continue
        if must_have_in_path is not None and must_have_in_path not in parts:
            continue

        owner_path = ".".join(parts[:-1])
        owner = modules.get(owner_path, root if owner_path == "" else None)
        if owner is None:
            continue
        candidates.append((owner, leaf, p, full_name))

    if not candidates:
        return None
    return candidates[0]


def transformers_gemma3_from_encoder(encoder):
    jsons_path = Path(__file__).parent / "gemma_configs"
    config = Gemma3Config.from_json_file(jsons_path / "gemma3cfg.json")
    with torch.device("meta"):
        metamodel = Gemma3ForConditionalGeneration(config)
    t_model_name = config.text_config.model_type
    t_model = _locate_model_within_model(metamodel, t_model_name)
    if t_model is None:
        raise ValueError(
            "Can't locate text model while converting text encoder to Gemma3ForConditionalGeneration"
        )
    t_model.load_state_dict(
        encoder.gemma3_12b.transformer.model.state_dict(), assign=True, strict=False
    )
    v_tower_name = config.vision_config.model_type
    v_tower = _locate_model_within_model(metamodel, v_tower_name)
    if v_tower is None:
        raise ValueError(
            "Can't locate vision model while converting text encoder to Gemma3ForConditionalGeneration"
        )
    v_model = v_tower.vision_model
    v_model.load_state_dict(
        encoder.gemma3_12b.transformer.vision_model.state_dict(),
        assign=True,
        strict=False,
    )
    metamodel.multi_modal_projector.load_state_dict(
        encoder.gemma3_12b.transformer.multi_modal_projector.state_dict(),
        assign=True,
        strict=False,
    )
    config = config.text_config
    dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
    base = config.rope_local_base_freq

    device = encoder.device
    positions_length = len(v_model.embeddings.position_ids[0])
    position_ids = torch.arange(
        positions_length, dtype=torch.long, device="cpu"
    ).unsqueeze(0)
    v_model.embeddings.register_buffer("position_ids", position_ids)
    embed_scale = torch.tensor(config.hidden_size**0.5, device=device)
    t_model.embed_tokens.register_buffer("embed_scale", embed_scale)
    local_rope_freqs = 1.0 / (
        base
        ** (
            torch.arange(0, dim, 2, dtype=torch.int64).to(
                device=device, dtype=torch.float
            )
            / dim
        )
    )
    t_model.rotary_emb_local.register_buffer("inv_freq", local_rope_freqs)
    rope_freqs, _ = ROPE_INIT_FUNCTIONS[config.rope_scaling["rope_type"]](
        config, device
    )
    t_model.rotary_emb.register_buffer("inv_freq", rope_freqs)
    lm_head_requires_grad = False
    loc_result = _locate_unique_parameter_owner_by_leaf(
        metamodel, leaf_param_name="weight", must_have_in_path="lm_head"
    )
    if loc_result is None:
        raise ValueError(
            "Can't locate lm_head while converting text encoder to Gemma3ForConditionalGeneration"
        )
    lm_head_owner, lm_head_attr, _, _ = loc_result
    real_w = t_model.embed_tokens.weight
    setattr(
        lm_head_owner,
        lm_head_attr,
        torch.nn.Parameter(real_w, requires_grad=lm_head_requires_grad),
    )
    metamodel.to(device)

    tokenizer_class = ltxv_gemma_tokenizer(jsons_path, max_length=1024)
    image_processor = AutoImageProcessor.from_pretrained(
        str(jsons_path),
        local_files_only=True,
    )
    processor = Gemma3Processor(
        image_processor=image_processor,
        tokenizer=tokenizer_class().tokenizer,
    )
    return metamodel, processor