import json as json_lib
import logging
import os
from collections.abc import Iterable
from typing import List, Optional, Set, Tuple

import torch
from torch import nn
from transformers import Llama4Config
from transformers.models.llama4.modeling_llama4 import (
    Llama4MultiModalProjector,
    Llama4VisionModel,
)

from sglang.srt.layers.logits_processor import LogitsProcessor
from sglang.srt.layers.moe.fused_moe_triton import FusedMoE
from sglang.srt.layers.quantization import QuantizationConfig
from sglang.srt.managers.mm_utils import (
    MultiModalityDataPaddingPatternMultimodalTokens,
    general_mm_embed_routine,
)
from sglang.srt.managers.schedule_batch import (
    Modality,
    MultimodalDataItem,
    MultimodalInputs,
)
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.model_loader.weight_utils import default_weight_loader
from sglang.srt.utils import add_prefix, is_cpu

_is_cpu = is_cpu()
from sglang.srt.model_loader.weight_utils import (
    default_weight_loader,
    maybe_remap_kv_scale_name,
)
from sglang.srt.utils import add_prefix

logger = logging.getLogger(__name__)


class Llama4ForConditionalGeneration(nn.Module):
    packed_modules_mapping = {
        "qkv_proj": ["q_proj", "k_proj", "v_proj"],
        "gate_up_proj": ["gate_proj", "up_proj"],
    }

    def __init__(
        self,
        config: Llama4Config,
        quant_config: Optional[QuantizationConfig] = None,
        prefix: str = "",
    ):
        super().__init__()
        self.config = config
        self.quant_config = quant_config

        # Check if this is a text-only model (modelopt fp8 llama4 has no vision components)
        self.has_vision = self._has_vision_weights(config)
        if not self.has_vision:
            logger.warning(
                "No vision weights found in checkpoint. Model will run in text-only mode. "
                "Multimodal capabilities (image processing) will be unavailable."
            )

        if self.has_vision:
            self.vision_model = Llama4VisionModel(config.vision_config)
            self.multi_modal_projector = Llama4MultiModalProjector(config)
        else:
            self.vision_model = None
            self.multi_modal_projector = None

        # Initialize the language model
        from sglang.srt.models.llama4 import Llama4ForCausalLM

        self.language_model = Llama4ForCausalLM(
            config.text_config if hasattr(config, "text_config") else config,
            quant_config=quant_config,
            prefix=add_prefix("language_model", prefix),
        )

        self.logits_processor = LogitsProcessor(
            config.text_config if hasattr(config, "text_config") else config
        )
        self.padding_pattern = MultiModalityDataPaddingPatternMultimodalTokens()

    def _has_vision_weights(self, config) -> bool:
        """Check if the model has vision components by examining the checkpoint."""
        model_path = getattr(config, "_name_or_path", None)
        if not model_path:
            return False

        # Check if this is a local path first
        if os.path.isdir(model_path):
            index_file = os.path.join(model_path, "model.safetensors.index.json")
            if os.path.exists(index_file):
                return self._check_vision_weights_in_index(index_file)

        # For HuggingFace models, we need to check the actual checkpoint
        # The config might say it's multimodal, but the checkpoint might be text-only
        try:
            # Try to access the HuggingFace cache directory
            from huggingface_hub import try_to_load_from_cache

            # Check if index file exists in cache
            index_file_path = try_to_load_from_cache(
                repo_id=model_path,
                filename="model.safetensors.index.json",
                cache_dir=None,
            )

            if index_file_path and os.path.exists(index_file_path):
                return self._check_vision_weights_in_index(index_file_path)

        except Exception:
            # If we can't access the cache, fall back to config-based detection
            pass

        # Fallback， assume text-only
        return False

    def _check_vision_weights_in_index(self, index_file: str) -> bool:
        """Check if the model.safetensors.index.json contains vision weights."""
        try:
            with open(index_file, "r") as f:
                index_data = json_lib.load(f)

            vision_patterns = ["vision_model", "vision_tower", "multi_modal_projector"]
            weight_names = index_data.get("weight_map", {}).keys()

            return any(
                pattern in weight_name
                for weight_name in weight_names
                for pattern in vision_patterns
            )
        except (OSError, json_lib.JSONDecodeError, KeyError):
            return False

    def pad_input_ids(self, input_ids: List[int], mm_inputs: MultimodalInputs):
        return self.padding_pattern.pad_input_tokens(input_ids, mm_inputs)

    def get_image_feature(
        self,
        items: List[MultimodalDataItem],
    ) -> torch.Tensor:
        # For text-only models, return None or raise an error
        if not self.has_vision or self.vision_model is None:
            raise ValueError("Vision model not available for text-only checkpoint")

        pixel_values = (
            torch.concat([item.feature for item in items])
            .to(next(self.vision_model.parameters()).device)
            .type(next(self.vision_model.parameters()).dtype)
        )

        image_outputs = self.vision_model(pixel_values, output_hidden_states=False)
        image_features = image_outputs.last_hidden_state
        vision_flat = image_features.view(-1, image_features.size(-1))
        projected_vision_flat = self.multi_modal_projector(vision_flat)
        return projected_vision_flat

    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
        forward_batch: ForwardBatch,
        **kwargs: object,
    ) -> torch.Tensor:

        # For text-only models, pass None for image_data_embedding_func
        image_embedding_func = self.get_image_feature if self.has_vision else None

        hs = general_mm_embed_routine(
            input_ids=input_ids,
            forward_batch=forward_batch,
            language_model=self.language_model,
            data_embedding_funcs={
                Modality.IMAGE: self.get_image_feature,
            },
            positions=positions,
        )

        return hs

    def permute_qk_weight_for_rotary(
        self,
        name: str,
        loaded_weight: torch.Tensor,
    ) -> Tuple[str, torch.Tensor]:

        def permute(w: torch.Tensor, n_heads: int):
            attn_in = self.language_model.config.head_dim * n_heads
            attn_out = self.language_model.config.hidden_size

            return (
                w.view(n_heads, attn_in // n_heads // 2, 2, attn_out)
                .transpose(1, 2)
                .reshape(attn_in, attn_out)
            )

        modules = name.split(".")

        # rotary embeds should be sliced
        if ("wk" in modules or "k_proj" in modules) and modules[-1] == "weight":
            if _is_cpu:
                dim = self.language_model.config.original_total_num_kv_heads
            else:
                dim = self.language_model.config.num_key_value_heads
            loaded_weight = permute(loaded_weight, dim)
        elif ("wq" in modules or "q_proj" in modules) and modules[-1] == "weight":
            if _is_cpu:
                dim = self.language_model.config.original_num_attention_heads
            else:
                dim = self.language_model.config.num_attention_heads
            loaded_weight = permute(loaded_weight, dim)

        return name, loaded_weight

    def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]) -> Set[str]:
        stacked_params_mapping = [
            # (param_name, shard_name, shard_id)
            (".self_attn.qkv_proj", ".self_attn.q_proj", "q"),
            (".self_attn.qkv_proj", ".self_attn.k_proj", "k"),
            (".self_attn.qkv_proj", ".self_attn.v_proj", "v"),
            (".shared_expert.gate_up_proj", ".shared_expert.gate_proj", 0),
            (".shared_expert.gate_up_proj", ".shared_expert.up_proj", 1),
            (".feed_forward.gate_up_proj", ".feed_forward.gate_proj", 0),
            (".feed_forward.gate_up_proj", ".feed_forward.up_proj", 1),
        ]

        params_dict = dict(self.named_parameters())
        num_experts = (
            self.config.text_config.num_local_experts
            if hasattr(self.config, "text_config")
            else self.config.num_local_experts
        )

        expert_params_mapping = FusedMoE.make_expert_params_mapping(
            ckpt_gate_proj_name="gate_proj",
            ckpt_down_proj_name="down_proj",
            ckpt_up_proj_name="up_proj",
            num_experts=num_experts,
        )

        for name, loaded_weight in weights:
            if self._should_skip_weight(name):
                continue

            name = self._transform_weight_name(name)

            if "vision" not in name:
                name, loaded_weight = self.permute_qk_weight_for_rotary(
                    name, loaded_weight
                )

            if self._handle_scale_remapping(name, params_dict):
                continue

            if self._handle_stacked_params(
                name, loaded_weight, stacked_params_mapping, params_dict
            ):
                continue

            if self._handle_expert_weights(
                name, loaded_weight, expert_params_mapping, params_dict, num_experts
            ):
                continue

            self._handle_default_weight(name, loaded_weight, params_dict)

    def _should_skip_weight(self, name: str) -> bool:
        """Check if we should skip loading this weight."""
        return "vision" in name and not self.has_vision

    def _transform_weight_name(self, name: str) -> str:
        """Transform weight name by adding language_model prefix if needed."""
        if (
            not name.startswith("language_model.")
            and "vision" not in name
            and "multi_modal_projector" not in name
        ):
            return f"language_model.{name}"
        return name

    def _handle_scale_remapping(self, name: str, params_dict: dict) -> bool:
        """Handle scale parameter remapping. Returns True if handled."""
        if "scale" in name and "expert" not in name:
            remapped_name = maybe_remap_kv_scale_name(name, params_dict)
            return remapped_name is None
        return False

    def _handle_stacked_params(
        self,
        name: str,
        loaded_weight: torch.Tensor,
        stacked_params_mapping: list,
        params_dict: dict,
    ) -> bool:
        """Handle stacked parameter loading. Returns True if handled."""
        for param_name, weight_name, shard_id in stacked_params_mapping:
            if weight_name in name and "vision" not in name:
                transformed_name = name.replace(weight_name, param_name)
                param = params_dict[transformed_name]
                param.weight_loader(param, loaded_weight, shard_id)
                return True
        return False

    def _handle_expert_weights(
        self,
        name: str,
        loaded_weight: torch.Tensor,
        expert_params_mapping: list,
        params_dict: dict,
        num_experts: int,
    ) -> bool:
        """Handle expert weight loading for MoE (Mixture of Experts) layers.

        Args:
            name: Parameter name from the checkpoint
            loaded_weight: The weight tensor to be loaded
            expert_params_mapping: Mapping of parameter names to expert configurations
            params_dict: Dictionary of model parameters
            num_experts: Total number of experts in the MoE layer

        Returns:
            bool: True if the parameter was handled (is an expert parameter), False otherwise
        """
        if ".experts" not in name:
            return False

        if "experts.gate_up_proj" not in name and "experts.down_proj" not in name:
            return self._handle_other_expert_params(
                name, loaded_weight, expert_params_mapping, params_dict
            )

        if "scale" in name:
            return self._handle_expert_scale_params(
                name, loaded_weight, params_dict, num_experts
            )
        else:
            return self._handle_expert_weight_params(
                name, loaded_weight, params_dict, num_experts
            )

    def _handle_other_expert_params(
        self,
        name: str,
        loaded_weight: torch.Tensor,
        expert_params_mapping: list,
        params_dict: dict,
    ) -> bool:
        """Handle expert parameters that are not gate_up_proj or down_proj weights.

        Args:
            name: Parameter name from the checkpoint
            loaded_weight: The weight tensor to be loaded
            expert_params_mapping: List of tuples mapping checkpoint names to model parameters
            params_dict: Dictionary of model parameters

        Returns:
            bool: True if parameter was found and handled, False otherwise
        """
        for param_name, weight_name, expert_id, shard_id in expert_params_mapping:
            if weight_name in name:
                transformed_name = name.replace(weight_name, param_name)
                param = params_dict[transformed_name]
                param.weight_loader(
                    param, loaded_weight, name, shard_id=shard_id, expert_id=expert_id
                )
                return True
        return False

    def _transform_expert_name(
        self, name: str, is_weight: bool = False
    ) -> Tuple[str, str, List[str]]:
        """Transform expert parameter name and get shard information.

        Args:
            name: The original parameter name
            is_weight: Whether this is a weight parameter (adds _weight suffix)

        Returns:
            Tuple of (transformed_name, shard_id, shard_id_list)
        """
        suffix = "_weight" if is_weight else ""

        if ".gate_up_proj" in name:
            transformed_name = name.replace(
                ".experts.gate_up_proj", f".experts.w13{suffix}"
            )
            shard_id = "w13"
            shard_id_list = ["w1", "w3"]
        else:  # down_proj
            transformed_name = name.replace(
                ".experts.down_proj", f".experts.w2{suffix}"
            )
            shard_id = "w2"
            shard_id_list = ["w2"]

        return transformed_name, shard_id, shard_id_list

    def _handle_expert_scale_params(
        self,
        name: str,
        loaded_weight: torch.Tensor,
        params_dict: dict,
        num_experts: int,
    ) -> bool:
        """Handle quantization scale parameters for expert weights.

        Args:
            name: Parameter name containing scale information
            loaded_weight: Scale tensor to be loaded
            params_dict: Dictionary of model parameters
            num_experts: Total number of experts for broadcast operations

        Returns:
            bool: True (always handles scale parameters)
        """
        import re

        # Check if this matches the expert parameter pattern: experts.{expert_id}.{param_name}
        expert_match = re.search(r"experts\.(\d+)\.", name)

        # Transform name
        transformed_name, _, _ = self._transform_expert_name(name)

        if transformed_name not in params_dict:
            return True

        param = params_dict[transformed_name]

        # Handle scale parameters
        if expert_match:
            # If we have a specific expert ID, only load for that expert
            expert_id = int(expert_match.group(1))
            # For scale parameters, we can directly set the value
            param.data[expert_id] = loaded_weight
        else:
            # No expert ID found - this is a single scale for all experts
            # Load the same scale for all experts
            for expert_id in range(num_experts):
                param.data[expert_id] = loaded_weight

        return True

    def _handle_expert_weight_params(
        self,
        name: str,
        loaded_weight: torch.Tensor,
        params_dict: dict,
        num_experts: int,
    ) -> bool:
        """Handle actual weight tensors for expert layers (gate_up_proj and down_proj).

        Args:
            name: Parameter name (should contain gate_up_proj or down_proj)
            loaded_weight: Weight tensor(s) to be loaded
            params_dict: Dictionary of model parameters
            num_experts: Total number of experts for tensor distribution

        Returns:
            bool: True (always handles weight parameters)
        """
        # Transform name and get shard info
        transformed_name, _, shard_id_list = self._transform_expert_name(
            name, is_weight=True
        )

        if ".gate_up_proj" in name:
            loaded_weight_list = loaded_weight.chunk(2, dim=-1)
        else:  # down_proj
            loaded_weight_list = [loaded_weight]

        for param_name, weight_chunk, shard_id in zip(
            [transformed_name] * len(shard_id_list), loaded_weight_list, shard_id_list
        ):
            if param_name not in params_dict:
                continue

            param = params_dict[param_name]
            weight_loader = param.weight_loader

            # Handle the case where loaded_weight might be a single tensor for all experts
            if weight_chunk.dim() == 2:
                # Single tensor case - load for all experts
                for expert_id in range(num_experts):
                    weight_loader(
                        param,
                        weight_chunk.T,
                        param_name,
                        shard_id=shard_id,
                        expert_id=expert_id,
                    )
            else:
                # Multiple experts case - load each expert's weights
                for expert_id in range(num_experts):
                    weight_loader(
                        param,
                        weight_chunk[expert_id].T,
                        param_name,
                        shard_id=shard_id,
                        expert_id=expert_id,
                    )

        return True

    def _handle_default_weight(
        self, name: str, loaded_weight: torch.Tensor, params_dict: dict
    ):
        """Handle default weight loading."""
        # Skip loading extra bias for GPTQ models
        if name.endswith(".bias") and name not in params_dict:
            return

        param = params_dict[name]
        weight_loader = getattr(param, "weight_loader", default_weight_loader)
        weight_loader(param, loaded_weight)

    def set_eagle3_layers_to_capture(self, layer_ids: Optional[List[int]] = None):
        if hasattr(self.language_model, "set_eagle3_layers_to_capture"):
            self.language_model.set_eagle3_layers_to_capture(layer_ids)

    def get_embed_and_head(self):
        # For EAGLE3, we delegate to the language model which should have this method
        # If the language model doesn't have lm_head (like EAGLE3), we return None for head
        embed = self.language_model.get_embed()
        if hasattr(self.language_model, "get_embed_and_head"):
            return self.language_model.get_embed_and_head()
        elif hasattr(self.language_model, "lm_head"):
            return embed, self.language_model.lm_head.weight
        else:
            # For EAGLE3, head might not be needed
            return embed, None

    def set_embed_and_head(self, embed, head):
        if hasattr(self.language_model, "set_embed_and_head"):
            return self.language_model.set_embed_and_head(embed, head)
        else:
            # For EAGLE3, only set embed
            return self.language_model.set_embed(embed)

    def get_embed(self):
        return self.language_model.get_embed()

    def set_embed(self, embed):
        return self.language_model.set_embed(embed)


EntryClass = Llama4ForConditionalGeneration
