"""
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

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

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

        Returns:
            A ModelGraph describing the full model structure.
        """
        local_path = self._resolve_model_path(model_path, hf_token)
        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)

        # Load model skeleton on meta device (zero memory)
        model_skeleton = self._load_skeleton(
            local_path, config, trust_remote_code
        )

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

        num_layers = len(layers)
        hidden_size = config.get(
            "hidden_size",
            config.get("d_model", config.get("n_embd", 0)),
        )
        vocab_size = config.get("vocab_size", 0)

        graph = ModelGraph(
            model_path=str(local_path),
            config=config,
            embedding=embedding,
            layers=layers,
            head=head,
            safetensors_index=safetensors_index,
            vocab_size=vocab_size,
            hidden_size=hidden_size,
            num_layers=num_layers,
        )

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

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

    def _resolve_model_path(
        self, model_path: str, hf_token: Optional[str]
    ) -> 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,
            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,
    ) -> nn.Module:
        """Load model architecture on meta device without allocating weights."""
        from accelerate import init_empty_weights
        from transformers import AutoConfig, AutoModelForCausalLM

        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

    # ------------------------------------------------------------------
    # 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}"

            # Classify this child
            if _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):
                # This is a MoE block — decompose into router + experts
                is_moe = True
                moe_components = self._decompose_moe(
                    full_child_name, child_module, param_meta, config
                )
                components.extend(moe_components)

                # Count experts
                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
                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
                    ))

            elif _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
                ))

            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."""
        components: list[ComponentInfo] = []

        for child_name, child_module in moe_module.named_children():
            full_name = f"{moe_name}.{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 isinstance(child_module, nn.ModuleList):
                # This is likely the experts list
                for j, expert in enumerate(child_module):
                    expert_name = f"{full_name}.{j}"
                    params = self._collect_params(expert_name, param_meta)
                    components.append(ComponentInfo(
                        name=expert_name, kind="expert", params=params
                    ))

            elif _EXPERT_PATTERNS.search(child_name):
                params = self._collect_params(full_name, param_meta)
                components.append(ComponentInfo(
                    name=full_name, 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 _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
        ]
