"""Model import: load PyTorch models and enumerate quantizable layers."""

from dataclasses import dataclass, field
from typing import Optional

import torch
import torch.nn as nn


@dataclass
class LayerInfo:
    name: str
    module_type: str  # "Linear" or "Conv2d"
    in_features: int
    out_features: int
    param_count: int
    weight_bytes: int  # FP16 size


@dataclass
class ModelInfo:
    model: object
    model_type: str  # "llm", "vlm", or "audio"
    layers: list[LayerInfo] = field(default_factory=list)
    total_params: int = 0
    total_weight_bytes: int = 0
    tokenizer: object = None
    model_name: str = ""


def _enumerate_linear_layers(model: nn.Module) -> list[LayerInfo]:
    layers = []
    for name, module in model.named_modules():
        if isinstance(module, nn.Linear):
            pc = module.weight.numel()
            if module.bias is not None:
                pc += module.bias.numel()
            layers.append(LayerInfo(
                name=name,
                module_type="Linear",
                in_features=module.in_features,
                out_features=module.out_features,
                param_count=pc,
                weight_bytes=pc * 2,  # FP16
            ))
    return layers


def _enumerate_conv_layers(model: nn.Module) -> list[LayerInfo]:
    layers = []
    for name, module in model.named_modules():
        if isinstance(module, (nn.Conv1d, nn.Conv2d)):
            pc = module.weight.numel()
            if module.bias is not None:
                pc += module.bias.numel()
            layers.append(LayerInfo(
                name=name,
                module_type=type(module).__name__,
                in_features=module.in_channels,
                out_features=module.out_channels,
                param_count=pc,
                weight_bytes=pc * 2,
            ))
    return layers


def load_hf_model(name: str, dtype: torch.dtype = torch.float16) -> ModelInfo:
    """Load a HuggingFace causal LM and enumerate quantizable layers."""
    from transformers import AutoModelForCausalLM, AutoTokenizer

    print(f"Loading HF model: {name}")
    tokenizer = AutoTokenizer.from_pretrained(name, trust_remote_code=True)
    model = AutoModelForCausalLM.from_pretrained(
        name,
        dtype=dtype,
        device_map="cpu",
        trust_remote_code=True,
    )
    model.eval()

    layers = _enumerate_linear_layers(model)
    total_params = sum(p.numel() for p in model.parameters())

    info = ModelInfo(
        model=model,
        model_type="llm",
        layers=layers,
        total_params=total_params,
        total_weight_bytes=total_params * 2,
        tokenizer=tokenizer,
        model_name=name,
    )
    print(f"  {total_params:,} params, {len(layers)} quantizable Linear layers")
    return info


def load_vlm_model(name: str, dtype: torch.dtype = torch.float16) -> ModelInfo:
    """Load a Vision-Language Model and enumerate quantizable layers.

    Supports models loadable via AutoModelForImageTextToText, AutoModel,
    etc. (e.g., Qwen3-VL, Qwen2-VL, SmolVLM, PaliGemma).
    """
    from transformers import AutoProcessor

    print(f"Loading VLM: {name}")

    # Try multiple loading strategies (order matters — most specific first)
    from transformers import AutoModel, AutoModelForCausalLM
    loader_classes = [
        ("AutoModelForImageTextToText", None),
        ("AutoModelForSeq2SeqLM", None),
        ("AutoModelForCausalLM", AutoModelForCausalLM),
        ("AutoModel", AutoModel),
    ]

    # Dynamically import optional classes
    try:
        from transformers import AutoModelForImageTextToText
        loader_classes[0] = ("AutoModelForImageTextToText", AutoModelForImageTextToText)
    except ImportError:
        pass
    try:
        from transformers import AutoModelForSeq2SeqLM
        loader_classes[1] = ("AutoModelForSeq2SeqLM", AutoModelForSeq2SeqLM)
    except ImportError:
        pass

    model = None
    last_error = None
    for cls_name, cls in loader_classes:
        if cls is None:
            continue
        try:
            model = cls.from_pretrained(
                name,
                dtype=dtype,
                device_map="cpu",
                trust_remote_code=True,
            )
            print(f"  Loaded via {cls_name}")
            break
        except Exception as e:
            last_error = e
            continue

    if model is None:
        raise RuntimeError(f"Could not load VLM: {name}. Last error: {last_error}")

    model.eval()
    processor = AutoProcessor.from_pretrained(name, trust_remote_code=True)

    layers = _enumerate_linear_layers(model)
    conv_layers = _enumerate_conv_layers(model)
    total_params = sum(p.numel() for p in model.parameters())

    # Categorize layers by component
    n_vision = sum(1 for l in layers if any(k in l.name for k in ["visual", "vision", "vit", "image_encoder", "img_encoder"]))
    n_proj = sum(1 for l in layers if any(k in l.name for k in ["multi_modal_projector", "merger", "connector", "proj"]))
    n_lm = len(layers) - n_vision - n_proj

    info = ModelInfo(
        model=model,
        model_type="vlm",
        layers=layers + conv_layers,
        total_params=total_params,
        total_weight_bytes=total_params * 2,
        tokenizer=processor,
        model_name=name,
    )
    print(f"  {total_params:,} params, {len(layers)} Linear layers "
          f"(~{n_vision} vision, ~{n_proj} projection, ~{n_lm} language)")
    print(f"  {len(conv_layers)} Conv2d layers")
    return info


def load_audio_model(name: str, dtype: torch.dtype = torch.float16) -> ModelInfo:
    """Load a Whisper speech recognition model and enumerate quantizable layers.

    Uses transformers WhisperForConditionalGeneration for PyTorch-side
    sensitivity analysis. MLX conversion uses mlx-whisper separately.
    """
    from transformers import WhisperForConditionalGeneration, WhisperProcessor

    print(f"Loading audio model: {name}")
    model = WhisperForConditionalGeneration.from_pretrained(
        name,
        torch_dtype=dtype,
        device_map="cpu",
    )
    model.eval()

    processor = WhisperProcessor.from_pretrained(name)

    layers = _enumerate_linear_layers(model)
    conv_layers = _enumerate_conv_layers(model)
    total_params = sum(p.numel() for p in model.parameters())

    # Categorize by encoder vs decoder
    # Note: decoder cross-attention layers have "encoder_attn" in the name
    # but live under "model.decoder.*", so check "model.encoder." prefix
    n_encoder = sum(1 for l in layers if l.name.startswith("model.encoder."))
    n_decoder = sum(1 for l in layers if l.name.startswith("model.decoder."))
    n_other = len(layers) - n_encoder - n_decoder

    info = ModelInfo(
        model=model,
        model_type="audio",
        layers=layers + conv_layers,
        total_params=total_params,
        total_weight_bytes=total_params * 2,
        tokenizer=processor,
        model_name=name,
    )
    print(f"  {total_params:,} params, {len(layers)} Linear layers "
          f"({n_encoder} encoder, {n_decoder} decoder)")
    print(f"  {len(conv_layers)} Conv1d/Conv2d layers")
    return info
