"""Standard HuggingFace loader mapping helpers."""

from __future__ import annotations

import functools
from collections.abc import Iterable, Sequence
from typing import Callable, Optional, Type  # noqa: UP035

import numpy as np
from tvm.relax.frontend import nn

from mlc_llm.loader import ExternMapping
from mlc_llm.quantization import Quantization

NameTransform = Callable[[str], str]
ExportSpecGetter = Callable[[nn.Module], object]


def _default_export_spec(model: nn.Module) -> object:
    return model.get_default_spec()


def make_standard_hf_loader(
    *,
    model_cls: Type[nn.Module],  # noqa: UP006
    layer_prefix: str = "model.layers",
    qkv_names: Sequence[str] = ("q_proj", "k_proj", "v_proj"),
    qkv_concat_axis: int = 0,
    qkv_target_name: str = "qkv_proj",
    add_qkv_bias: bool = False,
    qkv_bias_optional: bool = False,
    gate_up_names: Sequence[str] = ("gate_proj", "up_proj"),
    gate_up_concat_axis: int = 0,
    gate_up_target_name: str = "gate_up_proj",
    include_qkv: bool = True,
    include_gate_up: bool = True,
    add_unused: Optional[Iterable[str]] = None,  # noqa: UP045
    hf_prefix: str = "model.",
    name_transform: Optional[NameTransform] = None,  # noqa: UP045
    export_spec_getter: Optional[ExportSpecGetter] = None,  # noqa: UP045
    num_layers_getter: Optional[Callable[[object], int]] = None,  # noqa: UP045
) -> Callable[[object, Quantization], ExternMapping]:
    """Create a standard loader for HuggingFace weights.

    This handles the common QKV concatenation, gate+up concatenation, optional
    QKV bias mapping, and passes through remaining parameters 1:1.
    """

    if not qkv_names:
        include_qkv = False
    if not gate_up_names:
        include_gate_up = False
    if not include_qkv:
        qkv_names = ()
    if not include_gate_up:
        gate_up_names = ()

    def _default_name_transform(name: str) -> str:
        # When hf_prefix is empty, strip the "model." prefix so models that
        # expose bare top-level weights (no "model." namespace) still load.
        if hf_prefix == "":
            return name[6:] if name.startswith("model.") else name
        return name

    name_transform_fn = name_transform or _default_name_transform
    spec_getter = export_spec_getter or _default_export_spec
    unused_names = tuple(add_unused or ())

    def huggingface(
        model_config: object,
        quantization: Quantization,
    ) -> ExternMapping:
        model = model_cls(model_config)
        if quantization is not None:
            model.to(quantization.model_dtype)
        _, _named_params, _ = model.export_tvm(
            spec=spec_getter(model),
            allow_extern=True,
        )
        named_parameters = dict(_named_params)
        mapping = ExternMapping()

        if include_qkv or include_gate_up or unused_names:
            if num_layers_getter is None:
                num_layers = model_config.num_hidden_layers
            else:
                num_layers = num_layers_getter(model_config)

            for i in range(num_layers):
                attn = f"{layer_prefix}.{i}.self_attn"
                if include_qkv:
                    mlc_qkv_name = f"{attn}.{qkv_target_name}.weight"
                    mlc_param = named_parameters[mlc_qkv_name]
                    mapping.add_mapping(
                        mlc_qkv_name,
                        [name_transform_fn(f"{attn}.{name}.weight") for name in qkv_names],
                        functools.partial(
                            lambda q, k, v, dtype: np.concatenate(
                                [q, k, v], axis=qkv_concat_axis
                            ).astype(dtype),
                            dtype=mlc_param.dtype,
                        ),
                    )

                    if add_qkv_bias:
                        mlc_bias_name = f"{attn}.{qkv_target_name}.bias"
                        if (not qkv_bias_optional) or mlc_bias_name in named_parameters:
                            mlc_param = named_parameters[mlc_bias_name]
                            mapping.add_mapping(
                                mlc_bias_name,
                                [name_transform_fn(f"{attn}.{name}.bias") for name in qkv_names],
                                functools.partial(
                                    lambda q, k, v, dtype: np.concatenate(
                                        [q, k, v], axis=qkv_concat_axis
                                    ).astype(dtype),
                                    dtype=mlc_param.dtype,
                                ),
                            )

                if include_gate_up:
                    mlp = f"{layer_prefix}.{i}.mlp"
                    mlc_gate_up_name = f"{mlp}.{gate_up_target_name}.weight"
                    if gate_up_names:
                        mlc_param = named_parameters[mlc_gate_up_name]
                        mapping.add_mapping(
                            mlc_gate_up_name,
                            [name_transform_fn(f"{mlp}.{name}.weight") for name in gate_up_names],
                            functools.partial(
                                lambda gate, up, dtype: np.concatenate(
                                    [gate, up], axis=gate_up_concat_axis
                                ).astype(dtype),
                                dtype=mlc_param.dtype,
                            ),
                        )

                for unused_name in unused_names:
                    mapping.add_unused(name_transform_fn(f"{attn}.{unused_name}"))

        for mlc_name, mlc_param in named_parameters.items():
            if mlc_name not in mapping.param_map:
                mapping.add_mapping(
                    mlc_name,
                    [name_transform_fn(mlc_name)],
                    functools.partial(
                        lambda x, dtype: x.astype(dtype),
                        dtype=mlc_param.dtype,
                    ),
                )

        return mapping

    return huggingface