"""Utilities for navigating different HF model architectures."""

from __future__ import annotations

import torch.nn as nn

from obliteratus.models.loader import ModelHandle

# Mapping from model_type -> attribute path to the list of transformer layers.
_LAYER_ATTR_PATHS: dict[str, list[str]] = {
    "gpt2": ["transformer", "h"],
    "gpt_neo": ["transformer", "h"],
    "gpt_neox": ["gpt_neox", "layers"],
    "llama": ["model", "layers"],
    "mistral": ["model", "layers"],
    "mistral3": ["model", "language_model", "layers"],
    "mistral4": ["model", "layers"],
    "gemma": ["model", "layers"],
    "gemma2": ["model", "layers"],
    "phi": ["model", "layers"],
    "phi3": ["model", "layers"],
    "qwen2": ["model", "layers"],
    "qwen3": ["model", "layers"],
    "qwen3_moe": ["model", "layers"],
    "qwen3_5": ["model", "layers"],
    "qwen3_5_moe": ["model", "layers"],
    "qwen3_5_text": ["model", "layers"],
    "minimax_m2": ["model", "layers"],
    "glm_moe_dsa": ["model", "layers"],
    "deepseek_v3": ["model", "layers"],
    "glm4": ["model", "layers"],
    "glm4_moe": ["model", "layers"],
    "glm4_moe_lite": ["model", "layers"],
    "minicpm3": ["model", "layers"],
    "internlm3": ["model", "layers"],
    "falcon": ["transformer", "h"],
    "opt": ["model", "decoder", "layers"],
    "bloom": ["transformer", "h"],
    "mpt": ["transformer", "blocks"],
    "stablelm": ["model", "layers"],
    "chatglm": ["transformer", "encoder", "layers"],
    "glm": ["model", "layers"],
    "gpt_oss": ["model", "layers"],
    "smollm3": ["model", "layers"],
    "cohere": ["model", "layers"],
    "cohere2": ["model", "layers"],
    "olmo": ["model", "layers"],
    "olmo2": ["model", "layers"],
    "internlm2": ["model", "layers"],
    "granite": ["model", "layers"],
    "gemma3": ["model", "layers"],
    "nemotron_h": ["backbone", "layers"],
}

_LAYER_ATTR_PATH_FALLBACKS: dict[str, list[list[str]]] = {
    # Mistral3Model (without the conditional-generation head) exposes the
    # same text backbone one level closer to the root.
    "mistral3": [["language_model", "layers"]],
}

_ATTENTION_ATTR: dict[str, str] = {
    "gpt2": "attn",
    "gpt_neo": "attn.attention",
    "gpt_neox": "attention",
    "llama": "self_attn",
    "mistral": "self_attn",
    "mistral3": "self_attn",
    "mistral4": "self_attn",
    "gemma": "self_attn",
    "gemma2": "self_attn",
    "phi": "self_attn",
    "phi3": "self_attn",
    "qwen2": "self_attn",
    "qwen3": "self_attn",
    "qwen3_moe": "self_attn",
    "qwen3_5": "self_attn",
    "qwen3_5_moe": "self_attn",
    "qwen3_5_text": "self_attn",
    "minimax_m2": "self_attn",
    "glm_moe_dsa": "self_attn",
    "deepseek_v3": "self_attn",
    "glm4": "self_attn",
    "glm4_moe": "self_attn",
    "glm4_moe_lite": "self_attn",
    "minicpm3": "self_attn",
    "internlm3": "self_attn",
    "falcon": "self_attention",
    "opt": "self_attn",
    "bloom": "self_attention",
    "mpt": "attn",
    "stablelm": "self_attn",
    "chatglm": "self_attention",
    "glm": "self_attn",
    "gpt_oss": "self_attn",
    "smollm3": "self_attn",
    "cohere": "self_attn",
    "cohere2": "self_attn",
    "olmo": "self_attn",
    "olmo2": "self_attn",
    "internlm2": "attention",
    "granite": "self_attn",
    "gemma3": "self_attn",
    # Nemotron-H hybrid: attention (q/k/v/o_proj), Mamba (in/out_proj), and
    # MoE experts all live under layer.mixer.  Projection name matching
    # picks out whichever of those actually exists on the layer.
    "nemotron_h": "mixer",
}

_FFN_ATTR: dict[str, str] = {
    "gpt2": "mlp",
    "gpt_neo": "mlp",
    "gpt_neox": "mlp",
    "llama": "mlp",
    "mistral": "mlp",
    "mistral3": "mlp",
    "mistral4": "mlp",
    "gemma": "mlp",
    "gemma2": "mlp",
    "phi": "mlp",
    "phi3": "mlp",
    "qwen2": "mlp",
    "qwen3": "mlp",
    "qwen3_moe": "mlp",
    "qwen3_5": "mlp",
    "qwen3_5_moe": "mlp",
    "qwen3_5_text": "mlp",
    "minimax_m2": "mlp",
    "glm_moe_dsa": "mlp",
    "deepseek_v3": "mlp",
    "glm4": "mlp",
    "glm4_moe": "mlp",
    "glm4_moe_lite": "mlp",
    "minicpm3": "mlp",
    "internlm3": "mlp",
    "falcon": "mlp",
    # OPT: fc1/fc2 live directly on the layer — handled by _FLAT_FFN_ARCHS
    "bloom": "mlp",
    "mpt": "ffn",
    "stablelm": "mlp",
    "chatglm": "mlp",
    "glm": "mlp",
    "gpt_oss": "mlp",
    "smollm3": "mlp",
    "cohere": "mlp",
    "cohere2": "mlp",
    "olmo": "mlp",
    "olmo2": "mlp",
    "internlm2": "feed_forward",
    "granite": "mlp",
    "gemma3": "mlp",
    # Nemotron-H: MoE experts (mixer.experts) and shared experts
    # (mixer.shared_experts) hang off layer.mixer.
    "nemotron_h": "mixer",
}


# Architectures with hybrid attention (e.g. Qwen3.5 mixes standard multi-head
# attention with GatedDeltaNet).  On layers that lack the primary attribute,
# try the fallbacks in order.
_ATTENTION_ATTR_FALLBACKS: dict[str, list[str]] = {
    "qwen3_5": ["linear_attn"],
    "qwen3_5_moe": ["linear_attn"],
    "qwen3_5_text": ["linear_attn"],
}

# Architectures where FFN weights (fc1/fc2) live directly on the layer module
# instead of inside a dedicated MLP submodule.  For these, get_ffn_module
# returns the layer itself so _project_out_advanced can find fc1/fc2.
_FLAT_FFN_ARCHS: set[str] = {"opt"}


def _resolve_attr(obj, dotted_path: str):
    """Resolve a dotted attribute path like 'model.layers'."""
    for attr in dotted_path.split("."):
        obj = getattr(obj, attr)
    return obj


def get_layer_modules(handle: ModelHandle) -> nn.ModuleList:
    """Return the nn.ModuleList of transformer layers for this model."""
    arch = handle.architecture
    if arch in _LAYER_ATTR_PATHS:
        paths = [_LAYER_ATTR_PATHS[arch], *_LAYER_ATTR_PATH_FALLBACKS.get(arch, [])]
        for path in paths:
            obj = handle.model
            try:
                for attr in path:
                    obj = getattr(obj, attr)
            except AttributeError:
                continue
            if isinstance(obj, nn.ModuleList):
                return obj
        attempted = [".".join(path) for path in paths]
        raise RuntimeError(
            f"Cannot locate transformer layers for known architecture {arch!r}; "
            f"expected one of {attempted}."
        )

    # Fallback: walk the model looking for a ModuleList with the right length.
    # If num_layers is known, match exactly; otherwise find the largest ModuleList
    # (which is almost always the transformer layer stack).
    best = None
    for module in handle.model.modules():
        if isinstance(module, nn.ModuleList) and len(module) > 1:
            if handle.num_layers and len(module) == handle.num_layers:
                return module
            if best is None or len(module) > len(best):
                best = module
    if best is not None:
        return best
    raise RuntimeError(
        f"Cannot locate transformer layers for architecture {arch!r}. "
        f"Supported: {sorted(_LAYER_ATTR_PATHS)}"
    )


def get_attention_module(layer_module: nn.Module, architecture: str) -> nn.Module:
    """Return the attention sub-module inside a single transformer layer.

    For hybrid architectures (e.g. Qwen3.5 with GatedDeltaNet), tries
    fallback attribute names when the primary one is missing.
    """
    attr_path = _ATTENTION_ATTR.get(architecture, "self_attn")
    try:
        return _resolve_attr(layer_module, attr_path)
    except AttributeError:
        for fallback in _ATTENTION_ATTR_FALLBACKS.get(architecture, []):
            try:
                return _resolve_attr(layer_module, fallback)
            except AttributeError:
                continue
        raise


def get_ffn_module(layer_module: nn.Module, architecture: str) -> nn.Module:
    """Return the FFN/MLP sub-module inside a single transformer layer.

    For architectures with flat FFN layout (e.g. OPT where fc1/fc2 are
    direct attributes of the layer), returns the layer itself.
    """
    if architecture in _FLAT_FFN_ARCHS:
        return layer_module
    attr_path = _FFN_ATTR.get(architecture, "mlp")
    return _resolve_attr(layer_module, attr_path)


def get_embedding_module(handle: ModelHandle) -> nn.Embedding:
    """Return the token embedding module."""
    model = handle.model
    # Try common paths
    for path in [
        "transformer.wte",
        "model.language_model.embed_tokens",
        "language_model.embed_tokens",
        "model.embed_tokens",
        "gpt_neox.embed_in",
        "model.decoder.embed_tokens",
        "transformer.word_embeddings",
    ]:
        try:
            emb = _resolve_attr(model, path)
            if isinstance(emb, nn.Embedding):
                return emb
        except AttributeError:
            continue

    # Fallback: find first Embedding
    for module in model.modules():
        if isinstance(module, nn.Embedding):
            return module
    raise RuntimeError("Cannot locate embedding module.")
