"""
Model Graph Compiler - Model-agnostic computation graph analysis.

This is the core innovation that makes Dokodemo AI architecture-independent.
Instead of hardcoding support for specific model architectures (Llama, Mistral,
Mixtral, etc.), we automatically discover the model's structure by:

1. Loading the model skeleton on a meta device (zero memory)
2. Walking the module tree to find repeated layer blocks
3. Identifying MoE components (routers, experts) within layers
4. Mapping parameter names to SafeTensors file locations
5. Producing a universal ModelGraph for the inference engine
"""

from __future__ import annotations

import json
import logging
import os
import re
import struct
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Optional

from dokodemo_ai.graph.shape_inference import ShapeInference, ShapeInferredConfig

import torch
import torch.nn as nn
from huggingface_hub import snapshot_download

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Data structures representing the compiled model graph
# ---------------------------------------------------------------------------


@dataclass
class TensorRef:
    """Reference to a tensor stored in a SafeTensors file on disk."""

    name: str  # Full parameter name, e.g. "model.layers.0.self_attn.q_proj.weight"
    file: str  # SafeTensors filename (relative to model dir)
    dtype: str  # e.g. "F16", "BF16", "F32"
    shape: list[int] = field(default_factory=list)

    @property
    def numel(self) -> int:
        result = 1
        for s in self.shape:
            result *= s
        return result

    @property
    def size_bytes(self) -> int:
        bytes_per_elem = {"F32": 4, "F16": 2, "BF16": 2, "I8": 1, "I32": 4}
        return self.numel * bytes_per_elem.get(self.dtype, 2)


@dataclass
class ComponentInfo:
    """A sub-component within a layer (e.g. attention, MLP, router, expert)."""

    name: str  # Relative module name within the layer
    kind: str  # "attention", "mlp", "router", "expert", "norm", "other"
    params: list[TensorRef] = field(default_factory=list)

    @property
    def size_bytes(self) -> int:
        return sum(p.size_bytes for p in self.params)


@dataclass
class LayerInfo:
    """Complete description of one transformer layer."""

    index: int
    module_name: str  # e.g. "model.layers.5"
    is_moe: bool = False
    num_experts: int = 0
    num_active_experts: int = 0
    components: list[ComponentInfo] = field(default_factory=list)

    # Convenience accessors
    @property
    def attention(self) -> Optional[ComponentInfo]:
        return next((c for c in self.components if c.kind == "attention"), None)

    @property
    def router(self) -> Optional[ComponentInfo]:
        return next((c for c in self.components if c.kind == "router"), None)

    @property
    def experts(self) -> list[ComponentInfo]:
        return [c for c in self.components if c.kind == "expert"]

    @property
    def mlp(self) -> Optional[ComponentInfo]:
        return next((c for c in self.components if c.kind == "mlp"), None)

    @property
    def norms(self) -> list[ComponentInfo]:
        return [c for c in self.components if c.kind == "norm"]

    @property
    def all_params(self) -> list[TensorRef]:
        params = []
        for c in self.components:
            params.extend(c.params)
        return params

    @property
    def size_bytes(self) -> int:
        return sum(c.size_bytes for c in self.components)

    @property
    def expert_size_bytes(self) -> int:
        """Size of a single expert (average)."""
        experts = self.experts
        if not experts:
            return 0
        return sum(e.size_bytes for e in experts) // len(experts)

    @property
    def non_expert_size_bytes(self) -> int:
        """Size of everything except experts (attention + router + norms)."""
        return sum(
            c.size_bytes for c in self.components if c.kind != "expert"
        )


@dataclass
class EmbeddingInfo:
    """Embedding layer(s) at the start of the model."""

    module_name: str
    params: list[TensorRef] = field(default_factory=list)

    @property
    def size_bytes(self) -> int:
        return sum(p.size_bytes for p in self.params)


@dataclass
class HeadInfo:
    """Output head (LM head + final norm) at the end of the model."""

    module_name: str
    params: list[TensorRef] = field(default_factory=list)
    tied_to_embedding: bool = False

    @property
    def size_bytes(self) -> int:
        return sum(p.size_bytes for p in self.params)


@dataclass
class ModelGraph:
    """
    Complete, model-agnostic representation of a model's structure.

    This is the output of the ModelGraphCompiler and serves as the
    universal interface consumed by the inference engine, scheduler,
    cache, and streamer.
    """

    model_path: str
    config: dict[str, Any]
    embedding: EmbeddingInfo
    layers: list[LayerInfo]
    head: HeadInfo
    safetensors_index: dict[str, str]  # param_name -> filename
    vocab_size: int = 0
    hidden_size: int = 0
    num_layers: int = 0

    # VLM fields — populated only for vision-language models
    is_vlm: bool = False
    vision_arch: str = ""        # "clip" | "intern" | "qwen2" | ""
    vision_params: list = None   # list[str] — param names for the vision tower
    projector_params: list = None  # list[str] — param names for mm_projector

    # Weight-shape-inferred config (no config.json required)
    inferred: Optional[ShapeInferredConfig] = None

    @property
    def is_moe(self) -> bool:
        return any(layer.is_moe for layer in self.layers)

    @property
    def total_params(self) -> int:
        total = self.embedding.size_bytes
        for layer in self.layers:
            total += layer.size_bytes
        total += self.head.size_bytes
        return total

    @property
    def max_layer_size_bytes(self) -> int:
        if not self.layers:
            return 0
        return max(layer.size_bytes for layer in self.layers)

    def summary(self) -> str:
        lines = [
            f"Model: {self.model_path}",
            f"Architecture: {'MoE' if self.is_moe else 'Dense'}",
            f"Layers: {self.num_layers}",
            f"Hidden size: {self.hidden_size}",
            f"Vocab size: {self.vocab_size}",
            f"Embedding size: {self.embedding.size_bytes / 1e9:.2f} GB",
            f"Max layer size: {self.max_layer_size_bytes / 1e9:.2f} GB",
            f"Head size: {self.head.size_bytes / 1e9:.2f} GB",
        ]
        if self.is_moe:
            moe_layer = next(l for l in self.layers if l.is_moe)
            lines.extend([
                f"Experts per layer: {moe_layer.num_experts}",
                f"Active experts: {moe_layer.num_active_experts}",
                f"Expert size: {moe_layer.expert_size_bytes / 1e6:.1f} MB",
                f"Non-expert layer size: {moe_layer.non_expert_size_bytes / 1e6:.1f} MB",
                f"I/O savings from sparse loading: "
                f"{moe_layer.num_experts / max(moe_layer.num_active_experts, 1):.0f}x",
            ])
        return "\n".join(lines)


# ---------------------------------------------------------------------------
# Heuristic patterns for identifying model components
# ---------------------------------------------------------------------------

# Patterns that indicate a module is an attention block
_ATTN_PATTERNS = re.compile(
    r"(self_attn|attention|attn|self_attention)", re.IGNORECASE
)

# Patterns for MLP / feed-forward
_MLP_PATTERNS = re.compile(
    r"(^mlp$|feed_forward|ffn|dense_h_to)", re.IGNORECASE
)

# Patterns for MoE blocks
_MOE_PATTERNS = re.compile(
    r"(block_sparse_moe|moe|sparse_moe|mixture)", re.IGNORECASE
)

# Patterns for router / gating
_ROUTER_PATTERNS = re.compile(
    r"(gate$|router|gating|gate_proj_moe|w_gate)", re.IGNORECASE
)

# Patterns for expert modules
_EXPERT_PATTERNS = re.compile(
    r"(experts?\.\d+|expert_\d+|local_experts?\.\d+)", re.IGNORECASE
)

# Patterns for norm layers
_NORM_PATTERNS = re.compile(
    r"(layer_?norm|rms_?norm|input_layernorm|post_attention_layernorm|"
    r"pre_feedforward_layernorm|post_feedforward_layernorm)",
    re.IGNORECASE,
)

# Patterns for embedding
_EMBED_PATTERNS = re.compile(
    r"(embed_tokens|wte|word_embedding|tok_embed|embedding)", re.IGNORECASE
)

# Patterns for LM head
_HEAD_PATTERNS = re.compile(
    r"(lm_head|cls|output|score)", re.IGNORECASE
)

# Patterns for final norm
_FINAL_NORM_PATTERNS = re.compile(
    r"(model\.norm|model\.final_layernorm|transformer\.ln_f|"
    r"model\.layer_norm|gpt_neox\.final_layer_norm)",
    re.IGNORECASE,
)


# ---------------------------------------------------------------------------
# The Compiler
# ---------------------------------------------------------------------------


class ModelGraphCompiler:
    """
    Compiles any HuggingFace model into a universal ModelGraph.

    This is the key to model-agnosticism: instead of writing architecture-
    specific code, we analyze the model's module tree and parameter names
    to discover its structure automatically.
    """

    def compile(
        self,
        model_path: str,
        hf_token: Optional[str] = None,
        trust_remote_code: bool = False,
        cache_dir: Optional[str] = None,
    ) -> ModelGraph:
        """
        Compile a model from a HuggingFace repo ID or local path.

        Args:
            model_path: HuggingFace model ID or local directory path.
            hf_token: Optional HuggingFace API token for gated models.
            trust_remote_code: Whether to trust remote code in the model.
            cache_dir: Directory to cache downloaded models. Defaults to
                       HF_HOME env var, then ~/.cache/huggingface.

        Returns:
            A ModelGraph describing the full model structure.
        """
        local_path = self._resolve_model_path(model_path, hf_token, cache_dir=cache_dir)
        config = self._load_config(local_path)
        safetensors_index = self._load_safetensors_index(local_path)

        # Build param metadata from the index (no actual weight loading)
        param_meta = self._build_param_metadata(local_path, safetensors_index)

        # Infer architecture parameters from weight tensor shapes (config-free).
        # This supplements the HF config and enables loading models with incomplete
        # or absent config.json files.
        inferred = ShapeInference.from_index(safetensors_index, param_meta)

        # --- VLM detection (needed before skeleton load to pick right config) ---
        vision_params, projector_params, vision_arch = self._detect_vlm(
            safetensors_index, config
        )
        is_vlm = bool(vision_params)

        # Load model skeleton on meta device (zero memory).
        # For VLMs the text_config sub-config is used to get a plain LLM skeleton.
        model_skeleton, llm_prefix = self._load_skeleton(
            local_path, config, trust_remote_code, is_vlm=is_vlm
        )

        # For VLMs, safetensors LLM params have an extra prefix not present in
        # the text-only skeleton (e.g. "language_model." for LLaVA).  Build a
        # remapped param_meta that strips this prefix so _find_* discovery works
        # with the skeleton's module paths.  TensorRef.name retains the full
        # original name so weights are still loaded from the right file.
        disc_meta = param_meta  # alias; replaced below if VLM
        if llm_prefix:
            disc_meta = {
                name[len(llm_prefix):]: ref
                for name, ref in param_meta.items()
                if name.startswith(llm_prefix)
            }
            logger.info(
                "VLM: remapped %d LLM params, stripped prefix '%s'",
                len(disc_meta), llm_prefix,
            )

        # Discover model structure
        embedding = self._find_embedding(model_skeleton, disc_meta)
        layer_modules = self._find_layer_list(model_skeleton)
        layers = self._analyze_layers(
            layer_modules, disc_meta, config
        )
        head = self._find_head(
            model_skeleton, disc_meta, embedding
        )

        num_layers = len(layers)
        hidden_size = (
            config.get("hidden_size")
            or config.get("text_config", {}).get("hidden_size")
            or config.get("d_model", config.get("n_embd", 0))
            or inferred.hidden_size  # weight-derived fallback
        )
        vocab_size = config.get("vocab_size", 0) or inferred.vocab_size

        # Blend inferred values into config so InferenceEngine sees them.
        augmented_config = inferred.merge_with_config(config)

        # For VLMs (e.g. Kimi K2.5, LLaVA) the model-specific params like
        # kv_lora_rank, q_lora_rank, qk_nope_head_dim, etc. live inside the
        # text_config sub-dict.  Hoist them to the top level so that
        # InferenceEngine can find them with a simple config.get().
        tc = augmented_config.get("text_config")
        if isinstance(tc, dict):
            for k, v in tc.items():
                if k not in augmented_config and v is not None:
                    augmented_config[k] = v

        graph = ModelGraph(
            model_path=str(local_path),
            config=augmented_config,
            embedding=embedding,
            layers=layers,
            head=head,
            safetensors_index=safetensors_index,
            vocab_size=vocab_size,
            hidden_size=hidden_size,
            num_layers=num_layers,
            is_vlm=is_vlm,
            vision_arch=vision_arch,
            vision_params=vision_params or [],
            projector_params=projector_params or [],
            inferred=inferred,
        )

        logger.info("Compiled model graph:\n%s", graph.summary())
        return graph

    # ------------------------------------------------------------------
    # VLM detection
    # ------------------------------------------------------------------

    @staticmethod
    def _detect_vlm(
        safetensors_index: dict[str, str],
        config: dict,
    ) -> tuple[list[str], list[str], str]:
        """
        Detect vision-language model components in the safetensors index.

        Returns:
            (vision_params, projector_params, vision_arch)
            vision_params    — list of param names belonging to the vision tower
            projector_params — list of param names belonging to the mm_projector
            vision_arch      — "clip" | "intern" | "qwen2" | ""
        """
        # Vision tower prefixes per architecture
        VIT_PREFIXES = {
            "clip":   ["vision_tower.", "vision_model.", "model.vision_tower."],
            "intern": ["visual.", "model.visual.", "vision_model."],
            "qwen2":  ["visual.", "model.visual."],
        }
        # MLP projector prefixes
        PROJ_PREFIXES = [
            "mm_projector.", "model.mm_projector.", "projector.",
            "multi_modal_projector.",   # llava-hf format
        ]

        # Detect vision arch from model_type or config keys
        model_type = config.get("model_type", "").lower()
        arch = ""
        if "llava" in model_type or config.get("mm_vision_tower"):
            arch = "clip"
        elif "internvl" in model_type or "intern_vit" in model_type:
            arch = "intern"
        elif "qwen2_vl" in model_type or "qwen-vl" in model_type:
            arch = "qwen2"
        elif config.get("vision_config") and any(
            "vision_tower" in k or "vision_model" in k
            for k in safetensors_index
        ):
            arch = "clip"  # generic CLIP-style fallback

        if not arch:
            return [], [], ""

        # Collect vision params
        vision_prefixes = VIT_PREFIXES.get(arch, [])
        vision_params = [
            name for name in safetensors_index
            if any(name.startswith(p) for p in vision_prefixes)
            and not any(name.startswith(p) for p in PROJ_PREFIXES)
        ]
        projector_params = [
            name for name in safetensors_index
            if any(name.startswith(p) for p in PROJ_PREFIXES)
        ]

        if vision_params:
            vis_mb = sum(
                1 for n in vision_params  # rough count only, actual size computed at load
            )
            logger.info(
                "VLM detected: arch=%s, vision_tower=%d params, projector=%d params",
                arch, len(vision_params), len(projector_params),
            )

        return vision_params, projector_params, arch

    # ------------------------------------------------------------------
    # Path resolution
    # ------------------------------------------------------------------

    def _resolve_model_path(
        self,
        model_path: str,
        hf_token: Optional[str],
        cache_dir: Optional[str] = None,
    ) -> Path:
        """Download from HuggingFace Hub if needed, return local path."""
        local = Path(model_path)
        if local.is_dir() and (local / "config.json").exists():
            return local

        logger.info("Downloading model from HuggingFace Hub: %s", model_path)
        local_dir = snapshot_download(
            model_path,
            token=hf_token,
            cache_dir=cache_dir,          # None → uses HF_HOME env var
            ignore_patterns=["*.bin", "*.pt", "*.ckpt", "original/**"],
        )
        return Path(local_dir)

    # ------------------------------------------------------------------
    # Config and index loading
    # ------------------------------------------------------------------

    def _load_config(self, model_path: Path) -> dict:
        config_file = model_path / "config.json"
        if not config_file.exists():
            raise FileNotFoundError(f"No config.json found in {model_path}")
        with open(config_file) as f:
            return json.load(f)

    def _load_safetensors_index(self, model_path: Path) -> dict[str, str]:
        """
        Load the SafeTensors index mapping parameter names to shard files.

        Handles both single-file and multi-file (sharded) models.
        """
        index_file = model_path / "model.safetensors.index.json"
        if index_file.exists():
            with open(index_file) as f:
                data = json.load(f)
            return data.get("weight_map", {})

        # Single safetensors file
        single = model_path / "model.safetensors"
        if single.exists():
            from safetensors import safe_open

            index = {}
            with safe_open(str(single), framework="pt") as f:
                for key in f.keys():
                    index[key] = "model.safetensors"
            return index

        raise FileNotFoundError(
            f"No SafeTensors files found in {model_path}. "
            "Dokodemo AI requires models in SafeTensors format."
        )

    def _build_param_metadata(
        self, model_path: Path, index: dict[str, str]
    ) -> dict[str, TensorRef]:
        """
        Build metadata for all parameters without loading weights.

        Reads only the SafeTensors file header (a JSON blob at the start of
        every .safetensors file) to obtain dtype and shape for every tensor.
        This avoids using safe_open's PySafeSlice whose .dtype attribute was
        removed in newer safetensors versions.
        """
        meta: dict[str, TensorRef] = {}
        file_headers: dict[str, dict] = {}

        for param_name, filename in index.items():
            if filename not in file_headers:
                filepath = model_path / filename
                # SafeTensors format: first 8 bytes = little-endian uint64
                # header length, followed by that many bytes of UTF-8 JSON.
                with open(str(filepath), "rb") as f:
                    header_len = struct.unpack("<Q", f.read(8))[0]
                    header = json.loads(f.read(header_len))
                file_headers[filename] = header

            tensor_meta = file_headers[filename].get(param_name)
            if tensor_meta is None:
                # skip __metadata__ key or missing entries
                continue

            dtype = tensor_meta.get("dtype", "F16")   # e.g. "F16", "BF16"
            shape = tensor_meta.get("shape", [])

            meta[param_name] = TensorRef(
                name=param_name,
                file=filename,
                dtype=dtype,
                shape=shape,
            )

        return meta

    # ------------------------------------------------------------------
    # Model skeleton loading (meta device, zero memory)
    # ------------------------------------------------------------------

    def _load_skeleton(
        self,
        model_path: Path,
        config: dict,
        trust_remote_code: bool,
        is_vlm: bool = False,
    ) -> tuple[nn.Module, str]:
        """Load model architecture on meta device without allocating weights.

        For VLMs (LLaVA, InternVL2, Qwen2-VL) the HF model_type is a
        multimodal class not accepted by AutoModelForCausalLM.  We instead:
          1. Extract text_config and load a plain LLM skeleton from it.
          2. Compute the prefix that the safetensors index adds to skeleton
             param names (e.g. "language_model." for LLaVA).

        Returns:
            (skeleton, llm_prefix) — llm_prefix is "" for non-VLM models.
        """
        from accelerate import init_empty_weights
        from transformers import AutoConfig, AutoModelForCausalLM

        llm_prefix = ""

        if is_vlm and "text_config" in config:
            # Build an HF config from just the text sub-config
            text_cfg_dict = config["text_config"]
            # Inheriting the vocab_size from the outer config if missing
            if "vocab_size" not in text_cfg_dict and "vocab_size" in config:
                text_cfg_dict = {**text_cfg_dict, "vocab_size": config["vocab_size"]}
            _trust_remote = trust_remote_code
            try:
                text_hf_config = AutoConfig.for_model(**text_cfg_dict)
            except (ValueError, KeyError):
                # model_type is unknown to transformers (e.g. "kimi_k2" for
                # Kimi K2.5).  Load the full VLM config with trust_remote_code
                # so the custom config class is instantiated, then extract
                # its text sub-config (which is a proper PretrainedConfig).
                outer_cfg = AutoConfig.from_pretrained(
                    str(model_path), trust_remote_code=True
                )
                text_hf_config = getattr(outer_cfg, "text_config", outer_cfg)
                # Clear auto_map so from_config uses the standard registered
                # model class (avoids HFValidationError when _name_or_path is
                # empty and HF tries to load custom code from a hub repo_id).
                text_hf_config.auto_map = {}
                # Patch any attrs the standard model class expects that the
                # custom config may not have computed.
                # DeepSeek V3: qk_head_dim = qk_nope_head_dim + qk_rope_head_dim
                if (hasattr(text_hf_config, "qk_nope_head_dim")
                        and not hasattr(text_hf_config, "qk_head_dim")):
                    text_hf_config.qk_head_dim = (
                        text_hf_config.qk_nope_head_dim
                        + text_hf_config.qk_rope_head_dim
                    )
                # DeepSeek V3: transformers 4.57+ uses num_local_experts
                # but custom configs store it as n_routed_experts.
                if (hasattr(text_hf_config, "n_routed_experts")
                        and not hasattr(text_hf_config, "num_local_experts")):
                    text_hf_config.num_local_experts = (
                        text_hf_config.n_routed_experts
                    )
            with init_empty_weights():
                model = AutoModelForCausalLM.from_config(
                    text_hf_config, trust_remote_code=_trust_remote
                )

            # Detect the extra LLM prefix by matching a known skeleton param
            # against the safetensors index.  We look for "embed_tokens.weight"
            # in the skeleton and find its actual name in the safetensors index.
            llm_prefix = self._detect_llm_prefix(model, config)
            logger.info(
                "VLM skeleton loaded from text_config (%s).  LLM prefix: '%s'",
                text_cfg_dict.get("model_type", "?"), llm_prefix,
            )
            return model, llm_prefix

        # Standard text-only model
        hf_config = AutoConfig.from_pretrained(
            str(model_path), trust_remote_code=trust_remote_code
        )
        with init_empty_weights():
            model = AutoModelForCausalLM.from_config(
                hf_config, trust_remote_code=trust_remote_code
            )
        return model, llm_prefix

    @staticmethod
    def _detect_llm_prefix(skeleton: nn.Module, config: dict) -> str:
        """Find the prefix added to skeleton param names in the safetensors index.

        For LLaVA the skeleton (Llama) has 'model.embed_tokens.weight' but the
        safetensors file stores it as 'language_model.model.embed_tokens.weight'.
        The prefix is therefore 'language_model.'.

        We detect this by looking for known marker substrings in the config's
        safetensors_index (passed via config dict) or by inspecting the skeleton.
        """
        # Known VLM → LLM prefix mappings
        model_type = config.get("model_type", "").lower()
        if "llava" in model_type:
            return "language_model."
        if "internvl" in model_type:
            return "language_model."
        # Generic fallback: look for 'embed_tokens.weight' in skeleton params
        # and match against config architecture hints
        arch = (config.get("architectures") or [""])[0].lower()
        if "llava" in arch or "conditional" in arch:
            return "language_model."
        return ""

    # ------------------------------------------------------------------
    # Structural discovery
    # ------------------------------------------------------------------

    def _find_embedding(
        self,
        model: nn.Module,
        param_meta: dict[str, TensorRef],
    ) -> EmbeddingInfo:
        """Find the embedding layer(s)."""
        for name, module in model.named_modules():
            if isinstance(module, nn.Embedding) and _EMBED_PATTERNS.search(name):
                params = self._collect_params(name, param_meta)
                return EmbeddingInfo(module_name=name, params=params)

        # Fallback: first Embedding layer
        for name, module in model.named_modules():
            if isinstance(module, nn.Embedding):
                params = self._collect_params(name, param_meta)
                return EmbeddingInfo(module_name=name, params=params)

        raise RuntimeError("Could not find embedding layer in model")

    def _find_layer_list(
        self, model: nn.Module
    ) -> list[tuple[str, nn.Module]]:
        """
        Find the main repeated layer list (e.g. model.layers).

        We look for the largest nn.ModuleList whose children all have
        the same class — this is almost certainly the transformer layers.
        """
        candidates: list[tuple[str, nn.ModuleList]] = []

        for name, module in model.named_modules():
            if isinstance(module, nn.ModuleList) and len(module) > 1:
                # Check if children are homogeneous
                child_types = {type(c) for c in module}
                if len(child_types) == 1:
                    candidates.append((name, module))

        if not candidates:
            raise RuntimeError(
                "Could not find transformer layer list in model. "
                "Expected an nn.ModuleList with repeated layer blocks."
            )

        # Pick the longest ModuleList (the main layers)
        candidates.sort(key=lambda x: len(x[1]), reverse=True)
        list_name, layer_list = candidates[0]

        result = []
        for i, layer_module in enumerate(layer_list):
            full_name = f"{list_name}.{i}"
            result.append((full_name, layer_module))

        logger.info(
            "Found %d transformer layers at '%s'", len(result), list_name
        )
        return result

    def _analyze_layers(
        self,
        layer_modules: list[tuple[str, nn.Module]],
        param_meta: dict[str, TensorRef],
        config: dict,
    ) -> list[LayerInfo]:
        """Analyze each layer to identify its internal components."""
        layers = []
        for i, (layer_name, layer_module) in enumerate(layer_modules):
            layer_info = self._analyze_single_layer(
                i, layer_name, layer_module, param_meta, config
            )
            layers.append(layer_info)
        return layers

    def _analyze_single_layer(
        self,
        index: int,
        layer_name: str,
        layer_module: nn.Module,
        param_meta: dict[str, TensorRef],
        config: dict,
    ) -> LayerInfo:
        """Analyze a single transformer layer's internal structure."""
        components: list[ComponentInfo] = []
        is_moe = False
        num_experts = 0
        num_active_experts = 0

        # Get all direct children and their sub-modules
        for child_name, child_module in layer_module.named_children():
            full_child_name = f"{layer_name}.{child_name}"

            # NOTE: norm patterns are checked FIRST.  Names like
            # "post_attention_layernorm" contain "attention" and would otherwise
            # be misclassified as attention components.
            if _NORM_PATTERNS.search(child_name):
                params = self._collect_params(full_child_name, param_meta)
                components.append(ComponentInfo(
                    name=full_child_name, kind="norm", params=params
                ))

            elif _ATTN_PATTERNS.search(child_name):
                params = self._collect_params(full_child_name, param_meta)
                components.append(ComponentInfo(
                    name=full_child_name, kind="attention", params=params
                ))

            elif _MOE_PATTERNS.search(child_name):
                # Explicit MoE block (e.g., older Mixtral "block_sparse_moe")
                is_moe = True
                moe_components = self._decompose_moe(
                    full_child_name, child_module, param_meta, config
                )
                components.extend(moe_components)

                expert_comps = [c for c in moe_components if c.kind == "expert"]
                num_experts = len(expert_comps)
                num_active_experts = config.get(
                    "num_experts_per_tok",
                    config.get("num_selected_experts", config.get("top_k", 2)),
                )

            elif _MLP_PATTERNS.search(child_name):
                # Check if this MLP is actually a hidden MoE
                # (e.g., newer Mixtral/transformers renames block_sparse_moe → mlp)
                if self._has_expert_children(child_module):
                    is_moe = True
                    moe_components = self._decompose_moe(
                        full_child_name, child_module, param_meta, config
                    )
                    components.extend(moe_components)
                    expert_comps = [
                        c for c in moe_components if c.kind == "expert"
                    ]
                    num_experts = len(expert_comps)
                    num_active_experts = config.get(
                        "num_experts_per_tok", config.get("top_k", 2)
                    )
                else:
                    params = self._collect_params(full_child_name, param_meta)
                    components.append(ComponentInfo(
                        name=full_child_name, kind="mlp", params=params
                    ))

            else:
                # Unknown component — still track it
                params = self._collect_params(full_child_name, param_meta)
                if params:
                    components.append(ComponentInfo(
                        name=full_child_name, kind="other", params=params
                    ))

        # If we didn't detect MoE from children names, check config
        if not is_moe:
            num_experts_config = config.get(
                "num_local_experts", config.get("num_experts", 0)
            )
            if num_experts_config > 0:
                is_moe = True
                num_experts = num_experts_config
                num_active_experts = config.get(
                    "num_experts_per_tok",
                    config.get("top_k", 2),
                )

        return LayerInfo(
            index=index,
            module_name=layer_name,
            is_moe=is_moe,
            num_experts=num_experts,
            num_active_experts=num_active_experts,
            components=components,
        )

    def _decompose_moe(
        self,
        moe_name: str,
        moe_module: nn.Module,
        param_meta: dict[str, TensorRef],
        config: dict,
    ) -> list[ComponentInfo]:
        """Decompose a MoE block into router + individual experts.

        Handles two common issues:
        1. Python module path ≠ safetensors tensor path — e.g., newer
           transformers renamed Mixtral's ``block_sparse_moe`` attribute to
           ``mlp`` while the on-disk tensors kept the old name.
        2. Non-nn.ModuleList expert containers — e.g., ``MixtralExperts``
           (a custom class) instead of a plain ``nn.ModuleList``.
        """
        components: list[ComponentInfo] = []

        # Resolve the actual tensor prefix.  The Python attribute path (moe_name)
        # may not match the safetensors keys if HuggingFace renamed the module.
        actual_moe = self._find_actual_moe_name(moe_name, param_meta)

        for child_name, child_module in moe_module.named_children():
            full_name = f"{actual_moe}.{child_name}"

            if _ROUTER_PATTERNS.search(child_name):
                params = self._collect_params(full_name, param_meta)
                components.append(ComponentInfo(
                    name=full_name, kind="router", params=params
                ))

            elif "expert" in child_name.lower() or isinstance(child_module, nn.ModuleList):
                # Expert container — discover individual experts from param_meta
                # so we work with both nn.ModuleList and custom classes
                # (e.g., MixtralExperts in newer transformers).
                expert_comps = self._discover_experts_from_meta(full_name, param_meta)
                if expert_comps:
                    components.extend(expert_comps)
                elif isinstance(child_module, nn.ModuleList):
                    # Fallback to iterating the list if meta discovery fails
                    for j in range(len(child_module)):
                        ename = f"{full_name}.{j}"
                        params = self._collect_params(ename, param_meta)
                        components.append(ComponentInfo(
                            name=ename, kind="expert", params=params
                        ))

            elif _NORM_PATTERNS.search(child_name):
                params = self._collect_params(full_name, param_meta)
                components.append(ComponentInfo(
                    name=full_name, kind="norm", params=params
                ))

            else:
                params = self._collect_params(full_name, param_meta)
                if params:
                    components.append(ComponentInfo(
                        name=full_name, kind="other", params=params
                    ))

        return components

    def _find_actual_moe_name(
        self,
        python_moe_name: str,
        param_meta: dict[str, TensorRef],
    ) -> str:
        """Resolve the actual safetensors prefix for a MoE block.

        Transformers sometimes renames module attributes between releases
        while keeping the saved tensor names stable.  Example: Mixtral
        renamed ``block_sparse_moe`` → ``mlp`` in newer transformers, but
        the on-disk weights still use ``block_sparse_moe``.

        Strategy: if no params exist under ``python_moe_name``, scan
        siblings in the same layer for a component that looks like a MoE
        block (has ``gate``/``experts``/``router`` sub-components) and
        return that path instead.
        """
        if self._collect_params(python_moe_name, param_meta):
            return python_moe_name  # Fast path: name is already correct

        # Get parent prefix (e.g., "model.layers.0")
        parts = python_moe_name.rsplit(".", 1)
        if len(parts) != 2:
            return python_moe_name
        parent = parts[0]
        parent_prefix = parent + "."

        # Build map: child_name → set of its sub-component names
        child_subs: dict[str, set[str]] = {}
        for pname in param_meta:
            if pname.startswith(parent_prefix):
                rest = pname[len(parent_prefix):]
                segs = rest.split(".")
                if len(segs) >= 2:
                    child, sub = segs[0], segs[1]
                    child_subs.setdefault(child, set()).add(sub)

        # MoE-like children have gate/router/experts sub-components
        moe_indicators = {"gate", "router", "experts", "expert"}
        for child, subs in child_subs.items():
            actual = f"{parent}.{child}"
            if actual != python_moe_name and subs & moe_indicators:
                logger.info(
                    "MoE module path resolved: %r → %r (safetensors rename)",
                    python_moe_name, actual,
                )
                return actual

        return python_moe_name  # No alternative found; return original

    def _discover_experts_from_meta(
        self,
        experts_prefix: str,
        param_meta: dict[str, TensorRef],
    ) -> list[ComponentInfo]:
        """Discover individual expert components by scanning param_meta.

        Works for both ``nn.ModuleList`` experts and custom containers
        (e.g., ``MixtralExperts``), because it uses the saved tensor names
        rather than the Python module structure.
        """
        prefix = experts_prefix + "."
        expert_indices: set[int] = set()
        for pname in param_meta:
            if pname.startswith(prefix):
                rest = pname[len(prefix):]
                try:
                    idx = int(rest.split(".")[0])
                    expert_indices.add(idx)
                except ValueError:
                    pass

        return [
            ComponentInfo(
                name=f"{experts_prefix}.{j}",
                kind="expert",
                params=self._collect_params(f"{experts_prefix}.{j}", param_meta),
            )
            for j in sorted(expert_indices)
        ]

    def _has_expert_children(self, module: nn.Module) -> bool:
        """Check if a module contains expert sub-modules."""
        for name, _ in module.named_modules():
            if _EXPERT_PATTERNS.search(name) or _ROUTER_PATTERNS.search(name):
                return True
        return False

    def _find_head(
        self,
        model: nn.Module,
        param_meta: dict[str, TensorRef],
        embedding: EmbeddingInfo,
    ) -> HeadInfo:
        """Find the output head (LM head + final norm)."""
        head_params: list[TensorRef] = []
        head_name = ""
        tied = False

        # Look for LM head
        for name, module in model.named_modules():
            if _HEAD_PATTERNS.search(name) and isinstance(module, nn.Linear):
                head_name = name
                params = self._collect_params(name, param_meta)
                if params:
                    head_params.extend(params)
                else:
                    # Weight tying: LM head shares weights with embedding
                    tied = True
                    head_params = list(embedding.params)
                break

        # Look for final norm
        for name, module in model.named_modules():
            if _FINAL_NORM_PATTERNS.search(name):
                params = self._collect_params(name, param_meta)
                head_params.extend(params)
                if not head_name:
                    head_name = name
                break

        return HeadInfo(
            module_name=head_name,
            params=head_params,
            tied_to_embedding=tied,
        )

    # ------------------------------------------------------------------
    # Utility
    # ------------------------------------------------------------------

    def _collect_params(
        self, module_name: str, param_meta: dict[str, TensorRef]
    ) -> list[TensorRef]:
        """Collect all parameters whose names start with module_name."""
        prefix = module_name + "."
        return [
            ref
            for name, ref in param_meta.items()
            if name.startswith(prefix) or name == module_name
        ]
