"""
Text encoders for Module B (Context-aware Risk Detector).

Supports:
  - MacBERT-large (lightweight Chinese baseline)
  - Qwen2.5-7B with LoRA (full-scale Chinese)
  - LLaMA-3.1-8B with LoRA (multilingual)
"""

import torch
import torch.nn as nn
from transformers import AutoModel, AutoTokenizer
from peft import get_peft_model, LoraConfig, TaskType
from typing import Optional


class TextEncoder(nn.Module):
    """Shared backbone encoder for persona, context, and response."""

    def __init__(
        self,
        model_name: str,
        hidden_size: int = 768,
        use_lora: bool = False,
        lora_r: int = 16,
        lora_alpha: int = 32,
        lora_dropout: float = 0.05,
        freeze_base: bool = False,
    ):
        super().__init__()
        self.backbone = AutoModel.from_pretrained(model_name)
        self.actual_hidden = self.backbone.config.hidden_size

        if use_lora:
            lora_config = LoraConfig(
                task_type=TaskType.FEATURE_EXTRACTION,
                r=lora_r,
                lora_alpha=lora_alpha,
                lora_dropout=lora_dropout,
                target_modules=["q_proj", "v_proj", "query", "value"],
            )
            self.backbone = get_peft_model(self.backbone, lora_config)
        elif freeze_base:
            for param in self.backbone.parameters():
                param.requires_grad = False

        # Project to uniform hidden_size if needed
        self.proj = (
            nn.Linear(self.actual_hidden, hidden_size)
            if self.actual_hidden != hidden_size
            else nn.Identity()
        )

    def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
        """Returns [batch, seq_len, hidden_size]."""
        outputs = self.backbone(input_ids=input_ids, attention_mask=attention_mask)
        hidden = outputs.last_hidden_state
        return self.proj(hidden)

    def pool(self, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
        """Returns mean-pooled [batch, hidden_size]."""
        hidden = self.forward(input_ids, attention_mask)
        mask = attention_mask.unsqueeze(-1).float()
        return (hidden * mask).sum(1) / mask.sum(1).clamp(min=1e-9)


class ContextAwareFusion(nn.Module):
    """
    CrossAttention fusion: response as query, [persona; history] as key/value.
    Captures risk signals in response conditioned on relational context.
    """

    def __init__(self, hidden_size: int = 768, num_heads: int = 8, dropout: float = 0.1):
        super().__init__()
        self.cross_attn = nn.MultiheadAttention(
            embed_dim=hidden_size,
            num_heads=num_heads,
            dropout=dropout,
            batch_first=True,
        )
        self.layer_norm = nn.LayerNorm(hidden_size)
        self.ffn = nn.Sequential(
            nn.Linear(hidden_size, hidden_size * 4),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_size * 4, hidden_size),
        )
        self.ffn_norm = nn.LayerNorm(hidden_size)

    def forward(
        self,
        response_hidden: torch.Tensor,   # [B, R_len, H]
        context_hidden: torch.Tensor,    # [B, C_len, H]  (persona + history concat)
        context_key_padding_mask: Optional[torch.Tensor] = None,  # [B, C_len]
    ) -> torch.Tensor:
        """Returns [B, R_len, H] — response enriched with context signals."""
        attn_out, _ = self.cross_attn(
            query=response_hidden,
            key=context_hidden,
            value=context_hidden,
            key_padding_mask=context_key_padding_mask,
        )
        response_hidden = self.layer_norm(response_hidden + attn_out)
        ffn_out = self.ffn(response_hidden)
        return self.ffn_norm(response_hidden + ffn_out)
