"""Encoder stack used by the transformer architecture."""

from __future__ import annotations

import torch
import torch.nn as nn
import torch.utils.checkpoint as ckpt
from torch import Tensor

from .attention import MultiHeadAttention
from .feedforward import FeedForwardLayer

__all__ = ["EncoderLayer", "TransformerEncoder"]


class EncoderLayer(nn.Module):
    """Self-attention + feed-forward block supporting pre/post layer normalisation."""

    def __init__(
        self,
        d_model: int,
        num_heads: int,
        d_ff: int,
        dropout_rate: float,
        *,
        layer_norm_style: str = "post",
    ) -> None:
        super().__init__()

        if not isinstance(d_model, int):
            raise TypeError(f"d_model must be an int, got {type(d_model)}")
        if not isinstance(num_heads, int):
            raise TypeError(f"num_heads must be an int, got {type(num_heads)}")
        if not isinstance(d_ff, int):
            raise TypeError(f"d_ff must be an int, got {type(d_ff)}")
        if not isinstance(dropout_rate, float):
            raise TypeError(f"dropout_rate must be a float, got {type(dropout_rate)}")
        if not isinstance(layer_norm_style, str):
            raise TypeError(f"layer_norm_style must be a string, got {type(layer_norm_style)}")

        if d_model <= 0:
            raise ValueError("d_model must be strictly greater than 0")
        if num_heads <= 0:
            raise ValueError("num_heads must be strictly greater than 0")
        if d_ff <= 0:
            raise ValueError("d_ff must be strictly greater than 0")
        if not 0.0 <= dropout_rate < 1.0:
            raise ValueError("dropout_rate must be in [0, 1)")

        style = layer_norm_style.lower()
        if style not in {"pre", "post"}:
            raise ValueError("layer_norm_style must be either 'pre' or 'post' (case-insensitive)")
        self.layer_norm_style = style
        self.pre_norm = style == "pre"

        self.attention_layer = MultiHeadAttention(d_model, num_heads, dropout_rate)
        self.feed_forward = FeedForwardLayer(d_model, d_ff, dropout_rate)
        self.norm1 = nn.LayerNorm(d_model)
        self.dropout1 = nn.Dropout(dropout_rate)
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout2 = nn.Dropout(dropout_rate)

    def forward(self, x: Tensor, src_padding_mask: Tensor) -> Tensor:
        if not isinstance(x, torch.Tensor):
            raise TypeError("x must be a torch.Tensor")
        if x.dim() != 3:
            raise ValueError(
                f"x must be a 3D torch.Tensor of shape (B, S, D); got shape {tuple(x.shape)}"
            )
        if not isinstance(src_padding_mask, torch.Tensor):
            raise TypeError("src_padding_mask must be a torch.Tensor")
        if src_padding_mask.dtype != torch.bool or src_padding_mask.dim() != 4:
            raise TypeError(
                f"src_padding_mask must be boolean with shape (B, H, 1, S);"
                f" got dtype {src_padding_mask.dtype} and shape {tuple(src_padding_mask.shape)}"
            )

        if self.pre_norm:
            normed = self.norm1(x)
            attn_out = self.attention_layer(
                normed, normed, normed, src_padding_mask, src_padding_mask
            )
            x = x + self.dropout1(attn_out)
            ff_out = self.feed_forward(self.norm2(x))
            x = x + self.dropout2(ff_out)
            return x

        attn_out = self.attention_layer(x, x, x, src_padding_mask, src_padding_mask)
        x = self.norm1(x + self.dropout1(attn_out))
        ff_out = self.feed_forward(x)
        return self.norm2(x + self.dropout2(ff_out))


class TransformerEncoder(nn.Module):
    """Stack of encoder layers with optional activation checkpointing."""

    def __init__(
        self,
        d_model: int,
        num_heads: int,
        d_ff: int,
        num_layers: int,
        dropout_rate: float,
        *,
        layer_norm_style: str = "post",
    ) -> None:
        super().__init__()

        if not isinstance(num_layers, int):
            raise TypeError(f"num_layers must be an int, got {type(num_layers)}")
        if num_layers <= 0:
            raise ValueError("num_layers must be strictly greater than 0")

        if not isinstance(layer_norm_style, str):
            raise TypeError(f"layer_norm_style must be a string, got {type(layer_norm_style)}")
        style = layer_norm_style.lower()
        if style not in {"pre", "post"}:
            raise ValueError("layer_norm_style must be either 'pre' or 'post' (case-insensitive)")

        self.layer_norm_style = style
        self.layers = nn.ModuleList(
            [
                EncoderLayer(
                    d_model,
                    num_heads,
                    d_ff,
                    dropout_rate,
                    layer_norm_style=style,
                )
                for _ in range(num_layers)
            ]
        )
        self.use_ckpt = False

    def forward(self, x: Tensor, src_padding_mask: Tensor) -> Tensor:
        if not isinstance(x, torch.Tensor):
            raise TypeError("x must be a torch.Tensor")
        if x.dim() != 3:
            raise ValueError(
                f"x must be a 3D tensor of shape (B, S, D); got shape {tuple(x.shape)}"
            )

        for layer in self.layers:
            if self.use_ckpt:

                def _fn(x_, *, _layer=layer):
                    return _layer(x_, src_padding_mask)

                x = ckpt.checkpoint(_fn, x, use_reentrant=False)
            else:
                x = layer(x, src_padding_mask)
        return x