"""Compat patch for Lizzy TP under vLLM's generic Transformers backend."""

from __future__ import annotations

from typing import Any

import torch
import torch.nn.functional as F
from torch import nn

_PATCH_ATTR = "_flwr_transformers_lizzy_tp_patch_applied"


class _TensorParallelSliceNorm(nn.Module):
    """Apply a full-width checkpoint norm to a TP-local activation slice."""

    def __init__(self, base_norm: nn.Module, start_idx: int, end_idx: int):
        super().__init__()
        self.start_idx = start_idx
        self.end_idx = end_idx
        self.weight = base_norm.weight
        if getattr(base_norm, "bias", None) is not None:
            self.bias = base_norm.bias
        else:
            self.register_parameter("bias", None)
        self.eps = float(
            getattr(base_norm, "eps", getattr(base_norm, "variance_epsilon", 1e-6)),
        )
        self.norm_kind = (
            "layernorm" if isinstance(base_norm, nn.LayerNorm) else "rmsnorm"
        )

    @property
    def local_size(self) -> int:
        return self.end_idx - self.start_idx

    def _slice_param(self, param: torch.Tensor | None) -> torch.Tensor | None:
        if param is None:
            return None
        if param.shape[0] == self.local_size:
            return param
        return param[self.start_idx : self.end_idx]

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        weight = self.weight
        bias = self.bias
        if hidden_states.shape[-1] != self.weight.shape[0]:
            if hidden_states.shape[-1] != self.local_size:
                msg = (
                    "Unexpected hidden size for TP-sliced norm: "
                    f"{hidden_states.shape[-1]} "
                    f"(expected {self.weight.shape[0]} or {self.local_size})"
                )
                raise RuntimeError(msg)
            weight = self._slice_param(weight)
            bias = self._slice_param(bias)

        if self.norm_kind == "layernorm":
            return F.layer_norm(
                hidden_states,
                (hidden_states.shape[-1],),
                weight,
                bias,
                self.eps,
            )

        input_dtype = hidden_states.dtype
        hidden_states_fp32 = hidden_states.to(torch.float32)
        variance = hidden_states_fp32.pow(2).mean(dim=-1, keepdim=True)
        hidden_states_norm = hidden_states_fp32 * torch.rsqrt(variance + self.eps)
        hidden_states_norm = hidden_states_norm.to(input_dtype)
        output = weight * hidden_states_norm
        if bias is not None:
            output = output + bias
        return output


def _maybe_patch_lizzy_attention_for_tp(
    *,
    module: nn.Module,
    prefix: str,
    tp_size: int,
    tp_rank: int,
    log_replacement: Any,  # noqa: ANN401
) -> None:
    if tp_size <= 1 or type(module).__name__ != "LizzyAttention":
        return

    num_heads = getattr(module, "num_heads", None)
    num_key_value_heads = getattr(module, "num_key_value_heads", None)
    head_dim = getattr(module, "head_dim", None)
    q_norm = getattr(module, "q_norm", None)
    k_norm = getattr(module, "k_norm", None)
    if not all(
        isinstance(value, int)
        for value in (num_heads, num_key_value_heads, head_dim)
    ):
        return
    if num_heads % tp_size != 0 or num_key_value_heads % tp_size != 0:
        return

    local_num_heads = num_heads // tp_size
    local_num_key_value_heads = num_key_value_heads // tp_size
    local_q_dim = local_num_heads * head_dim
    local_kv_dim = local_num_key_value_heads * head_dim

    module.num_heads = local_num_heads
    module.num_key_value_heads = local_num_key_value_heads
    module.num_key_value_groups = local_num_heads // local_num_key_value_heads

    if q_norm is not None and getattr(q_norm, "weight", None) is not None:
        start = tp_rank * local_q_dim
        end = start + local_q_dim
        module.q_norm = _TensorParallelSliceNorm(q_norm, start, end)
        log_replacement(f"{prefix}.q_norm", q_norm, module.q_norm)

    if k_norm is not None and getattr(k_norm, "weight", None) is not None:
        start = tp_rank * local_kv_dim
        end = start + local_kv_dim
        module.k_norm = _TensorParallelSliceNorm(k_norm, start, end)
        log_replacement(f"{prefix}.k_norm", k_norm, module.k_norm)


def patch_vllm_transformers_lizzy_tp() -> None:
    """Patch the generic vLLM Transformers backend for Lizzy TP norms/heads."""
    import vllm.model_executor.models.transformers as transformers_mod

    transformers_base = transformers_mod.TransformersBase
    if getattr(transformers_base, _PATCH_ATTR, False):
        return

    PreTrainedModel = transformers_mod.PreTrainedModel
    maybe_prefix = transformers_mod.maybe_prefix
    replace_linear_class = transformers_mod.replace_linear_class
    get_feature_request_tip = transformers_mod.get_feature_request_tip
    re = transformers_mod.re
    log_replacement = transformers_mod.log_replacement
    get_tp_rank = getattr(transformers_mod, "get_tensor_model_parallel_rank", None)
    if get_tp_rank is None:
        try:
            from vllm.distributed import (  # noqa: PLC0415
                get_tensor_model_parallel_rank as get_tp_rank,
            )
        except Exception:
            get_tp_rank = lambda: 0

    def tensor_parallel(self: Any) -> None:  # noqa: ANN401
        """Apply the model's tensor parallel plan plus Lizzy attention fixes."""
        is_pretrained_model = lambda m: isinstance(m, PreTrainedModel)
        supports_tp_plan = lambda m: m.config.base_model_tp_plan is not None
        pretrained_models = filter(is_pretrained_model, self.model.modules())
        models_with_tp_plan = filter(supports_tp_plan, pretrained_models)

        if not any(models_with_tp_plan) and self.tp_size > 1:
            tip = get_feature_request_tip(
                self.model_config.model,
                self.model_config.trust_remote_code,
            )
            raise ValueError(
                f"{type(self.model)} does not support tensor parallel. {tip}",
            )

        tp_rank = get_tp_rank()

        def _tensor_parallel(
            module: nn.Module,
            prefix: str = "",
            tp_plan: dict[str, str] | None = None,
        ) -> None:
            local_tp_plan = tp_plan or {}

            if isinstance(module, PreTrainedModel):
                local_tp_plan = module.config.base_model_tp_plan or {}
                local_tp_plan = {
                    maybe_prefix(prefix, key): value
                    for key, value in local_tp_plan.items()
                }

            for child_name, child_module in module.named_children():
                qual_name = maybe_prefix(prefix, child_name)
                if isinstance(child_module, nn.Linear):
                    generator = (p for p in local_tp_plan if re.match(p, qual_name))
                    pattern = next(generator, None)
                    style = local_tp_plan.get(pattern, "replicate")
                    new_module = replace_linear_class(
                        child_module,
                        style,
                        self.quant_config,
                        prefix=qual_name,
                    )
                    setattr(module, child_name, new_module)
                    log_replacement(qual_name, child_module, new_module)
                else:
                    _tensor_parallel(
                        child_module,
                        prefix=qual_name,
                        tp_plan=local_tp_plan,
                    )

            _maybe_patch_lizzy_attention_for_tp(
                module=module,
                prefix=prefix,
                tp_size=self.tp_size,
                tp_rank=tp_rank,
                log_replacement=log_replacement,
            )

        _tensor_parallel(self.model)

    transformers_base.tensor_parallel = tensor_parallel
    setattr(transformers_base, _PATCH_ATTR, True)
