"""Multi-head attention building block with explicit validation helpers."""

from __future__ import annotations

from typing import cast

import torch
import torch.nn as nn
from torch import Tensor

from ..utils import (
    calculate_attention,
    combine_masks,
    create_qk_padding_mask,
    join_heads,
    split_heads,
)

__all__ = ["MultiHeadAttention"]


class MultiHeadAttention(nn.Module):
    """
    Multi-head attention: linear projections -> split heads -> scaled dot-product
    attention (via utils) -> merge heads -> output projection (+ dropout).

    Args:
        d_model (int): Model dimension (>0). Must be divisible by num_heads.
        num_heads (int): Number of attention heads (>0).
        dropout_rate (float): Dropout probability in (0,1).

    Inputs:
        query, key, value: (B, S, D) with D == d_model
        mask (optional): Boolean tensor broadcastable to (B, H, S_q, S_k) where True entries are masked.

    Returns:
        Tensor: (B, S_q, D)
    """

    def __init__(self, d_model: int, num_heads: int, dropout_rate: float):
        super().__init__()
        # ---- type checks
        if not isinstance(num_heads, int):
            raise TypeError(f"num_heads must be an int, got {type(num_heads)}")
        if not isinstance(d_model, int):
            raise TypeError(f"d_model must be an int, got {type(d_model)}")
        if not isinstance(dropout_rate, float):
            raise TypeError(f"dropout_rate must be a float, got {type(dropout_rate)}")

        # ---- value checks
        if num_heads <= 0:
            raise ValueError(f"num_heads must be strictly greater than 0, got {num_heads}")
        if d_model <= 0:
            raise ValueError(f"d_model must be strictly greater than 0, got {d_model}")
        if d_model % num_heads != 0:
            raise ValueError("d_model must be divisible by num_heads")
        if not (0 <= dropout_rate < 1):
            raise ValueError(f"dropout_rate must be between 0 and 1 excluded, got {dropout_rate}")

        self.d_model = d_model
        self.num_heads = num_heads
        self.d_head = d_model // num_heads
        self.dropout_rate = dropout_rate

        self.query_linear = nn.Linear(d_model, d_model, bias=False)
        self.key_linear = nn.Linear(d_model, d_model, bias=False)
        self.value_linear = nn.Linear(d_model, d_model, bias=False)
        self.output_linear = nn.Linear(d_model, d_model, bias=True)

    def forward(
        self,
        query: Tensor,
        key: Tensor,
        value: Tensor,
        q_mask: Tensor,
        k_mask: Tensor,
        causal_mask: Tensor | None = None,
    ) -> Tensor:
        """Run multi-head attention with boolean padding and optional causal masks."""

        # ---- basic type checks
        if not isinstance(query, torch.Tensor):
            raise TypeError(f"query must be a torch.Tensor, got {type(query)}")
        if not isinstance(key, torch.Tensor):
            raise TypeError(f"key must be a torch.Tensor, got {type(key)}")
        if not isinstance(value, torch.Tensor):
            raise TypeError(f"value must be a torch.Tensor, got {type(value)}")
        if not isinstance(q_mask, torch.Tensor):
            raise TypeError(f"q_mask must be a torch.Tensor, got {type(q_mask)}")
        if k_mask is None or not isinstance(k_mask, torch.Tensor):
            raise TypeError(f"k_mask must be a torch.Tensor, got {type(k_mask)}")

        if query.dim() != 3 or key.dim() != 3 or value.dim() != 3:
            raise ValueError(
                "query/key/value must be 3D tensors of shape (B, S, D); "
                f"got q={tuple(query.shape)}, k={tuple(key.shape)}, v={tuple(value.shape)}"
            )
        Bq, Sq, Dq = query.shape
        Bk, Sk, Dk = key.shape
        Bv, Sv, Dv = value.shape

        if not (Dq == Dk == Dv == self.d_model):
            raise ValueError(
                f"Last dimension must equal d_model={self.d_model}; got Dq={Dq}, Dk={Dk}, Dv={Dv}"
            )
        if not (Bq == Bk == Bv):
            raise ValueError(f"Batch size mismatch: q={Bq}, k={Bk}, v={Bv}")
        if Sk != Sv:
            raise ValueError(f"Key/Value seq length mismatch: Sk={Sk} vs Sv={Sv}")

        # ---- padding mask validation ----
        if q_mask.dtype != torch.bool:
            raise TypeError(f"q_mask must be boolean, got {q_mask.dtype}")
        if k_mask.dtype != torch.bool:
            raise TypeError(f"k_mask must be boolean, got {k_mask.dtype}")
        if q_mask.dim() != 4 or k_mask.dim() != 4:
            raise ValueError(
                "q_mask and k_mask must be 4D tensors shaped (B, H, 1, S); "
                f"got {tuple(q_mask.shape)} and {tuple(k_mask.shape)}"
            )

        # ---- project and split into heads -> (B, H, S, Dh)
        q = split_heads(self.query_linear(query), self.num_heads)
        k = split_heads(self.key_linear(key), self.num_heads)
        v = split_heads(self.value_linear(value), self.num_heads)
        pad_mask = create_qk_padding_mask(q_mask, k_mask)
        combined_mask = combine_masks(causal_mask, pad_mask)

        p = self.dropout_rate if self.training else 0.0
        attn = cast(Tensor, calculate_attention(q, k, v, combined_mask, dropout_p=p))
        out = join_heads(attn)  # (B, Sq, D)
        return self.output_linear(out)