"""Decoder stack used by the encoder-decoder transformer."""

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__ = ["DecoderLayer", "TransformerDecoder"]


class DecoderLayer(nn.Module):
    """Single decoder block with self-attention, cross-attention, and feed-forward."""

    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.self_attention_layer = MultiHeadAttention(d_model, num_heads, dropout_rate)
        self.cross_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)
        self.norm3 = nn.LayerNorm(d_model)
        self.dropout3 = nn.Dropout(dropout_rate)

    def forward(
        self,
        x: Tensor,
        y: Tensor,
        src_padding_mask: Tensor,
        tgt_padding_mask: Tensor,
        tgt_causal_mask: Tensor | None,
    ) -> Tensor:
        """Run one decoder layer using the configured layer-normalisation style."""

        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(y, torch.Tensor):
            raise TypeError("y must be a torch.Tensor")
        if y.dim() != 3:
            raise ValueError(
                f"y must be a 3D torch.Tensor of shape (B, S, D); got shape {tuple(y.shape)}"
            )
        if x.shape[0] != y.shape[0] or x.shape[-1] != y.shape[-1]:
            raise ValueError("Encoder memory and decoder input must match in batch and d_model")

        for mask_name, mask in (
            ("src_padding_mask", src_padding_mask),
            ("tgt_padding_mask", tgt_padding_mask),
        ):
            if not isinstance(mask, torch.Tensor):
                raise TypeError(f"{mask_name} must be a torch.Tensor")
            if mask.dtype != torch.bool or mask.dim() != 4:
                raise TypeError(
                    f"{mask_name} must be boolean with shape (B, H, 1, S);"
                    f" got dtype {mask.dtype} and shape {tuple(mask.shape)}"
                )

        if self.pre_norm:
            normed = self.norm1(y)
            self_attn_out = self.self_attention_layer(
                normed,
                normed,
                normed,
                tgt_padding_mask,
                tgt_padding_mask,
                tgt_causal_mask,
            )
            y = y + self.dropout1(self_attn_out)

            normed_y = self.norm2(y)
            cross_attn_out = self.cross_attention_layer(
                normed_y,
                x,
                x,
                tgt_padding_mask,
                src_padding_mask,
            )
            y = y + self.dropout2(cross_attn_out)

            ff_out = self.feed_forward(self.norm3(y))
            y = y + self.dropout3(ff_out)
            return y

        # Self-attention (post-LN)
        self_attn_out = self.self_attention_layer(
            y,
            y,
            y,
            tgt_padding_mask,
            tgt_padding_mask,
            tgt_causal_mask,
        )
        y = self.norm1(y + self.dropout1(self_attn_out))

        # Cross-attention (encoder memory as keys/values)
        cross_attn_out = self.cross_attention_layer(
            y,
            x,
            x,
            tgt_padding_mask,
            src_padding_mask,
        )
        y = self.norm2(y + self.dropout2(cross_attn_out))

        # Feed-forward block
        ff_out = self.feed_forward(y)
        return self.norm3(y + self.dropout3(ff_out))


class TransformerDecoder(nn.Module):
    """Stack of decoder 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(
            [
                DecoderLayer(
                    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,
        y: Tensor,
        src_padding_mask: Tensor,
        tgt_padding_mask: Tensor,
        tgt_causal_mask: Tensor | None,
    ) -> Tensor:
        """Run the decoder stack for all time steps."""

        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(y, torch.Tensor):
            raise TypeError("y must be a torch.Tensor")
        if y.dim() != 3:
            raise ValueError(
                f"y must be a 3D torch.Tensor of shape (B, S, D); got shape {tuple(y.shape)}"
            )
        if x.shape[0] != y.shape[0] or x.shape[-1] != y.shape[-1]:
            raise ValueError("Encoder memory and decoder input must match in batch and d_model")

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

                def _fn(y_, *, _layer=layer):
                    return _layer(x, y_, src_padding_mask, tgt_padding_mask, tgt_causal_mask)

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