import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import AutoModel, AutoTokenizer
from typing import List, Literal, Optional

PoolingType = Literal["mean", "cls", "max"]


class MizanEmbeddingModel(nn.Module):
    """Base Mizan embedding model (v1).

    - backbone_name: HuggingFace transformer model id
    - emb_dim: output embedding dimension
    - pooling: 'mean', 'cls', or 'max'
    - normalize: whether to L2-normalize embeddings (usually False for Mizan)
    """

    def __init__(
        self,
        backbone_name: str = "distilbert-base-uncased",
        emb_dim: int = 384,
        pooling: PoolingType = "mean",
        normalize: bool = False,
    ):
        super().__init__()
        self.backbone_name = backbone_name
        self.backbone = AutoModel.from_pretrained(backbone_name)
        hidden = self.backbone.config.hidden_size
        self.proj = nn.Linear(hidden, emb_dim)
        self.pooling: PoolingType = pooling
        self.normalize = normalize

    def _pool(self, token_embeddings: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
        """Pool token embeddings into a single vector per sequence."""

        if self.pooling == "cls":
            # Use first token (e.g., [CLS])
            return token_embeddings[:, 0, :]

        mask = attention_mask.unsqueeze(-1).float()

        if self.pooling == "mean":
            summed = (token_embeddings * mask).sum(dim=1)       # [B, H]
            counts = mask.sum(dim=1).clamp(min=1e-8)            # [B, 1]
            return summed / counts

        if self.pooling == "max":
            replaced = token_embeddings.masked_fill(mask == 0, -1e9)
            return replaced.max(dim=1).values

        raise ValueError(f"Unknown pooling: {self.pooling}")

    def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
        outputs = self.backbone(input_ids=input_ids, attention_mask=attention_mask)
        token_embeddings = outputs.last_hidden_state  # [B, L, H]
        pooled = self._pool(token_embeddings, attention_mask)   # [B, H]
        emb = self.proj(pooled)                                 # [B, D]
        if self.normalize:
            emb = F.normalize(emb, p=2, dim=-1)
        return emb


class MizanTextEncoderWrapper:
    """Convenience wrapper around MizanEmbeddingModel for text inference.

    Example:
        encoder = MizanTextEncoderWrapper()
        emb = encoder.encode_one("some text")
    """

    def __init__(
        self,
        backbone_name: str = "distilbert-base-uncased",
        emb_dim: int = 384,
        pooling: PoolingType = "mean",
        normalize: bool = False,
        device: Optional[str] = None,
    ):
        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
        self.tokenizer = AutoTokenizer.from_pretrained(backbone_name)
        self.model = MizanEmbeddingModel(
            backbone_name=backbone_name,
            emb_dim=emb_dim,
            pooling=pooling,
            normalize=normalize,
        ).to(self.device)
        self.model.eval()

    @torch.no_grad()
    def encode(self, texts: List[str], batch_size: int = 16) -> torch.Tensor:
        """Encode a list of texts into embeddings [N, D]."""

        all_embs = []
        for i in range(0, len(texts), batch_size):
            batch = texts[i : i + batch_size]
            enc = self.tokenizer(
                batch,
                padding=True,
                truncation=True,
                max_length=256,
                return_tensors="pt",
            )
            enc = {k: v.to(self.device) for k, v in enc.items()}
            emb = self.model(**enc).cpu()
            all_embs.append(emb)
        return torch.cat(all_embs, dim=0)

    @torch.no_grad()
    def encode_one(self, text: str) -> torch.Tensor:
        """Encode a single text into a 1D embedding tensor [D]."""

        return self.encode([text])[0]
