"""
Transformer Decoder 模块 — Person B 负责实现

包含:
- TransformerDecoderLayer: 单层解码器
- TransformerDecoder: 多层解码器堆叠

架构 (Pre-LayerNorm):
    x → LN → Masked Self-Attention → Residual
      → LN → Cross-Attention → Residual
      → LN → FFN → Residual
"""

from __future__ import annotations

import torch
import torch.nn as nn
from typing import Optional


class TransformerDecoderLayer(nn.Module):
    """
    单层 Transformer Decoder。

    TODO [Person B]: 实现以下组件:
    1. Masked Self-Attention (因果掩码，防止看到未来)
    2. Cross-Attention (decoder 查询 encoder 输出)
    3. Feed-Forward Network
    4. 三个 LayerNorm
    5. Residual connections + Dropout
    """

    def __init__(
        self,
        d_model: int = 512,
        nhead: int = 8,
        dim_feedforward: int = 2048,
        dropout: float = 0.1,
        activation: str = "gelu",
        use_flash_attention: bool = True,
        use_rotary_embedding: bool = True,
        pre_norm: bool = True,
    ):
        super().__init__()
        raise NotImplementedError("TODO: Person B 实现 TransformerDecoderLayer.__init__")

    def forward(
        self,
        tgt: torch.Tensor,                               # [B, T, D]
        memory: torch.Tensor,                             # [B, S, D] (encoder output)
        tgt_mask: Optional[torch.Tensor] = None,          # [T, T] causal mask
        memory_key_padding_mask: Optional[torch.BoolTensor] = None,  # [B, S]
        tgt_key_padding_mask: Optional[torch.BoolTensor] = None,     # [B, T]
    ) -> torch.Tensor:
        """
        TODO [Person B]: Pre-LayerNorm 前向传播:
        1. Masked Self-Attention with causal mask
        2. Cross-Attention with encoder output
        3. FFN
        每步都有 residual connection 和 dropout
        """
        raise NotImplementedError("TODO: Person B 实现 TransformerDecoderLayer.forward")


class TransformerDecoder(nn.Module):
    """
    多层 Transformer Decoder。

    TODO [Person B]:
    1. 堆叠 N 个 TransformerDecoderLayer
    2. 最终加一个 LayerNorm
    """

    def __init__(self, decoder_layer: TransformerDecoderLayer, num_layers: int):
        super().__init__()
        raise NotImplementedError("TODO: Person B 实现 TransformerDecoder.__init__")

    def forward(
        self,
        tgt: torch.Tensor,
        memory: torch.Tensor,
        tgt_mask: Optional[torch.Tensor] = None,
        memory_key_padding_mask: Optional[torch.BoolTensor] = None,
        tgt_key_padding_mask: Optional[torch.BoolTensor] = None,
    ) -> torch.Tensor:
        raise NotImplementedError("TODO: Person B 实现 TransformerDecoder.forward")