"""Stream transformer with KV caching for SalTTS."""

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


class StreamTransformer(nn.Module):
    """Causal transformer with streaming attention and KV caching.

    Features:
    - 12 layers, 512d hidden, 8 attention heads
    - KV cache for efficient streaming (512 tokens)
    - Cross-attention (Q from text1, K/V from text2)
    - Chunked processing (1-5 tokens, lookahead 3)
    """

    def __init__(
        self,
        num_layers: int = 12,
        hidden_dim: int = 512,
        num_heads: int = 8,
        ff_dim: int = 2048,
        dropout: float = 0.1,
        max_seq_len: int = 512,
        kv_cache_size: int = 512,
    ):
        """Initialize stream transformer.

        Args:
            num_layers: Number of transformer layers
            hidden_dim: Hidden dimension
            num_heads: Number of attention heads
            ff_dim: Feed-forward dimension
            dropout: Dropout probability
            max_seq_len: Maximum sequence length
            kv_cache_size: KV cache size
        """
        super().__init__()

        self.hidden_dim = hidden_dim
        self.num_heads = num_heads
        self.kv_cache_size = kv_cache_size

        # Embedding layer
        self.embedding = nn.Embedding(1000, hidden_dim)  # Vocab size placeholder
        self.pos_encoding = nn.Parameter(torch.randn(1, max_seq_len, hidden_dim))

        # Transformer layers
        self.layers = nn.ModuleList([
            TransformerLayer(hidden_dim, num_heads, ff_dim, dropout)
            for _ in range(num_layers)
        ])

        # Cross-attention layers
        self.cross_attn_layers = nn.ModuleList([
            CrossAttentionLayer(hidden_dim, num_heads, dropout)
            for _ in range(num_layers)
        ])

        # Output projection
        self.output_proj = nn.Linear(hidden_dim, hidden_dim)

        # KV cache
        self.kv_cache = None

    def forward(
        self,
        text1_ids: torch.Tensor,
        text2_emb: Optional[torch.Tensor] = None,
        use_cache: bool = True,
    ) -> torch.Tensor:
        """Forward pass.

        Args:
            text1_ids: Token IDs from text1 [batch, seq_len]
            text2_emb: Embeddings from text2 [batch, seq_len2, 512]
            use_cache: Whether to use KV caching

        Returns:
            Output features [batch, seq_len, 512]
        """
        batch_size, seq_len = text1_ids.shape

        # Embed and add positional encoding
        x = self.embedding(text1_ids)  # [batch, seq_len, 512]
        x = x + self.pos_encoding[:, :seq_len, :]

        # Process through transformer layers
        for i, (layer, cross_attn) in enumerate(zip(self.layers, self.cross_attn_layers)):
            # Self-attention with causal mask
            x = layer(x, use_cache=use_cache)

            # Cross-attention with text2 if provided
            if text2_emb is not None:
                x = cross_attn(x, text2_emb)

        # Output projection
        output = self.output_proj(x)  # [batch, seq_len, 512]

        return output

    def reset_cache(self) -> None:
        """Reset KV cache."""
        self.kv_cache = None


class TransformerLayer(nn.Module):
    """Single transformer layer with causal self-attention."""

    def __init__(self, hidden_dim: int, num_heads: int, ff_dim: int, dropout: float):
        super().__init__()

        self.self_attn = nn.MultiheadAttention(hidden_dim, num_heads, dropout=dropout, batch_first=True)
        self.norm1 = nn.LayerNorm(hidden_dim)

        self.ffn = nn.Sequential(
            nn.Linear(hidden_dim, ff_dim),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(ff_dim, hidden_dim),
            nn.Dropout(dropout),
        )
        self.norm2 = nn.LayerNorm(hidden_dim)

    def forward(self, x: torch.Tensor, use_cache: bool = True) -> torch.Tensor:
        # Self-attention with residual
        attn_out, _ = self.self_attn(x, x, x, need_weights=False)
        x = self.norm1(x + attn_out)

        # Feed-forward with residual
        ffn_out = self.ffn(x)
        x = self.norm2(x + ffn_out)

        return x


class CrossAttentionLayer(nn.Module):
    """Cross-attention layer (Q from text1, K/V from text2)."""

    def __init__(self, hidden_dim: int, num_heads: int, dropout: float):
        super().__init__()

        self.cross_attn = nn.MultiheadAttention(hidden_dim, num_heads, dropout=dropout, batch_first=True)
        self.norm = nn.LayerNorm(hidden_dim)

    def forward(self, query: torch.Tensor, key_value: torch.Tensor) -> torch.Tensor:
        # Cross-attention with residual
        attn_out, _ = self.cross_attn(query, key_value, key_value, need_weights=False)
        output = self.norm(query + attn_out)

        return output
