import gguf
import numpy as np
from constants import ARCH_FEATURES


def detect_architecture(tensors: dict) -> str:
    """Detect model architecture from tensor names."""
    names = list(tensors.keys())

    has_ssm = any("ssm_" in n for n in names)
    has_qkv = any("attn_qkv" in n for n in names)
    has_nextn = any("nextn" in n for n in names)
    has_moe = any("exps" in n for n in names)
    has_separate_qkv = any("attn_q.weight" in n for n in names)
    has_gemma_specific = any(
        t in n for t in ("layer_output_scale", "post_attention_norm", "post_ffw_norm")
        for n in names
    )

    if has_ssm and has_qkv:
        return "qwen35"
    if has_moe:
        return "mellum2"
    # Исправление: заменяем неопределённую has_blk_attn на уже существующий признак
    if has_separate_qkv or has_gemma_specific:
        return "gemma4"
    return "unknown"


def _detect_prefix(tensors: dict) -> str:
    for name in tensors:
        if name.startswith("BLK."):
            return "BLK"
    return "blk"


def _estimate_layers(tensors: dict) -> int:
    max_layer = 0
    for name in tensors:
        parts = name.split(".")
        if len(parts) >= 2 and parts[0] in ("blk", "BLK"):
            try:
                layer = int(parts[1])
                if layer > max_layer:
                    max_layer = layer
            except ValueError:
                pass
    return max_layer + 1  # layers are 0-indexed


def read_model(path: str) -> dict:
    """Parse BF16 GGUF, return model info."""
    r = gguf.GGUFReader(path)
    tensors = {}
    meta = {}

    for k, v in r.fields.items():
        # Безопасное извлечение данных (список/число, а не raw numpy)
        try:
            data = v.data
            if isinstance(data, np.ndarray):
                data = data.tolist()
            elif isinstance(data, (np.generic,)):
                data = data.item()
            meta[k] = data
        except Exception:
            meta[k] = str(v)

    for t in r.tensors:
        shape = list(t.shape)
        name = t.name
        n_elements = int(np.prod(shape))
        tensors[name] = {
            "shape": shape,
            "n_elements": n_elements,
            "size_mib": n_elements * 2 / 1024 / 1024,
        }

    arch = detect_architecture(tensors)
    arch_features = ARCH_FEATURES.get(arch, {}).copy()

    prefix = _detect_prefix(tensors)
    n_layers = _estimate_layers(tensors)

    if arch == "mellum2" and arch_features.get("moe_intermediate_size", 0) == 0:
        arch_features["moe_intermediate_size"] = 896

    arch_features["prefix"] = prefix
    if n_layers > 0:
        arch_features["n_layers"] = n_layers

    has_moe = any("exps" in n for n in tensors)
    if has_moe:
        arch_features["has_moe"] = True

    has_nextn = any("nextn" in n for n in tensors)
    # blk.32 is MTP only in ~32-layer models (Qwen). Skip for deeper models (Gemma4, 48 layers).
    has_blk32 = n_layers <= 33 and any(
        n.startswith("blk.32.") or n.startswith("BLK.32.") for n in tensors
    )
    arch_features["has_mtp"] = has_nextn or has_blk32

    return {
        "path": path,
        "architecture": arch,
        "features": arch_features,
        "tensors": tensors,
        "n_tensors": len(tensors),
        "meta": meta,
    }