"""
HuggingFace PreTrainedModel for mentee-embed.

Supports:
  AutoModel.from_pretrained("MenteEAI/mentee-embed-v3", trust_remote_code=True)

The model returns last_hidden_state + pooler_output (mean-pooled, L2-normalised)
so it also works as a drop-in with sentence-transformers >= 2.2 via:
  SentenceTransformer("MenteEAI/mentee-embed-v3", trust_remote_code=True)
"""
from __future__ import annotations

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import PreTrainedModel
from transformers.modeling_outputs import BaseModelOutputWithPooling

from .configuration_mentee import MenteeEmbedConfig


# ── building blocks ──────────────────────────────────────────────────────────

class _Block(nn.Module):
    def __init__(self, cfg: MenteeEmbedConfig):
        super().__init__()
        self.ln1 = nn.LayerNorm(cfg.hidden)
        self.attn = nn.MultiheadAttention(
            cfg.hidden, cfg.heads, dropout=cfg.dropout, batch_first=True
        )
        self.ln2 = nn.LayerNorm(cfg.hidden)
        self.ffn = nn.Sequential(
            nn.Linear(cfg.hidden, cfg.ffn),
            nn.GELU(),
            nn.Linear(cfg.ffn, cfg.hidden),
            nn.Dropout(cfg.dropout),
        )
        self.drop = nn.Dropout(cfg.dropout)

    def forward(self, x: torch.Tensor, key_padding_mask: torch.Tensor) -> torch.Tensor:
        h = self.ln1(x)
        attn_out, _ = self.attn(h, h, h, key_padding_mask=key_padding_mask, need_weights=False)
        x = x + self.drop(attn_out)
        x = x + self.ffn(self.ln2(x))
        return x


class _TextEncoder(nn.Module):
    def __init__(self, cfg: MenteeEmbedConfig):
        super().__init__()
        self.tok_emb = nn.Embedding(cfg.vocab_size, cfg.hidden, padding_idx=0)
        self.pos_emb = nn.Embedding(cfg.max_position, cfg.hidden)
        self.drop = nn.Dropout(cfg.dropout)
        self.blocks = nn.ModuleList([_Block(cfg) for _ in range(cfg.layers)])
        self.ln_f = nn.LayerNorm(cfg.hidden)

    def forward(self, input_ids: torch.Tensor):
        B, L = input_ids.shape
        pos = torch.arange(L, device=input_ids.device).unsqueeze(0).expand(B, L)
        x = self.tok_emb(input_ids) + self.pos_emb(pos)
        x = self.drop(x)
        pad_mask = input_ids.eq(0)
        for blk in self.blocks:
            x = blk(x, pad_mask)
        return self.ln_f(x)                # (B, L, H)  — last_hidden_state


# ── main model ───────────────────────────────────────────────────────────────

class MenteeEmbedModel(PreTrainedModel):
    """
    mentee-embed encoder.

    Returns BaseModelOutputWithPooling:
      .last_hidden_state  — (B, L, H)  token-level representations
      .pooler_output      — (B, H)     mean-pooled, L2-normalised sentence embedding

    Quick usage:
        from transformers import AutoModel, AutoTokenizer
        model = AutoModel.from_pretrained("MenteEAI/mentee-embed-v3", trust_remote_code=True)
        # use model.encode(texts) for the simplest path
    """

    config_class = MenteeEmbedConfig
    base_model_prefix = "encoder"

    def __init__(self, config: MenteeEmbedConfig):
        super().__init__(config)
        self.encoder = _TextEncoder(config)
        h = config.hidden
        self.proj = nn.Sequential(
            nn.Linear(h, h),
            nn.GELU(),
            nn.Linear(h, h),
        )
        self.post_init()

    # ── core forward (HF-standard signature) ─────────────────────────────────
    def forward(
        self,
        input_ids: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
        **kwargs,
    ) -> BaseModelOutputWithPooling:
        last_hidden = self.encoder(input_ids)           # (B, L, H)

        # mean-pool over non-padding tokens
        if attention_mask is not None:
            mask = attention_mask.unsqueeze(-1).to(last_hidden.dtype)
        else:
            mask = (~input_ids.eq(0)).unsqueeze(-1).to(last_hidden.dtype)

        pooled = (last_hidden * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1.0)
        pooled = self.proj(pooled)                      # projection head
        pooled = F.normalize(pooled, dim=-1)            # L2-normalise

        return BaseModelOutputWithPooling(
            last_hidden_state=last_hidden,
            pooler_output=pooled,
        )

    # ── convenience helper ───────────────────────────────────────────────────
    @torch.no_grad()
    def encode(
        self,
        texts: list[str],
        tokenizer,
        batch_size: int = 64,
        device: str | None = None,
    ) -> torch.Tensor:
        """
        Encode a list of strings into L2-normalised embeddings.

        Args:
            texts:      list of strings to encode
            tokenizer:  the tokenizer returned by MenteeTokenizer.from_pretrained(...)
                        OR any callable that returns {"input_ids": tensor}
            batch_size: inference batch size
            device:     "cpu" / "cuda" / None (auto-detect)

        Returns:
            FloatTensor of shape (len(texts), hidden_dim), L2-normalised
        """
        if device is None:
            device = "cuda" if torch.cuda.is_available() else "cpu"
        self.to(device).eval()

        all_embs = []
        for i in range(0, len(texts), batch_size):
            batch = texts[i : i + batch_size]
            enc = tokenizer(batch, return_tensors="pt", padding=True, truncation=True, max_length=512)
            input_ids = enc["input_ids"].to(device)
            attention_mask = enc.get("attention_mask")
            if attention_mask is not None:
                attention_mask = attention_mask.to(device)
            out = self(input_ids, attention_mask=attention_mask)
            all_embs.append(out.pooler_output.cpu())

        return torch.cat(all_embs, dim=0)