import os
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"]


# ============================================================
#  Embedding Model
# ============================================================
class MizanEmbeddingModel(nn.Module):
    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 = pooling
        self.normalize = normalize

    def _pool(self, token_embeddings, attention_mask):
        if self.pooling == "cls":
            return token_embeddings[:, 0, :]

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

        if self.pooling == "mean":
            summed = (token_embeddings * mask).sum(dim=1)
            counts = mask.sum(dim=1).clamp(min=1e-8)
            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, attention_mask, token_type_ids=None):
        outputs = self.backbone(
            input_ids=input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
        )
        token_embeddings = outputs.last_hidden_state
        pooled = self._pool(token_embeddings, attention_mask)
        emb = self.proj(pooled)
        if self.normalize:
            emb = F.normalize(emb, p=2, dim=-1)
        return emb

    # ------------------------------------------------------------
    # Save model
    # ------------------------------------------------------------
    def save_pretrained(self, save_directory: str):
        os.makedirs(save_directory, exist_ok=True)

        torch.save(self.state_dict(), os.path.join(save_directory, "pytorch_model.bin"))

        config = {
            "backbone_name": self.backbone_name,
            "emb_dim": self.proj.out_features,
            "pooling": self.pooling,
            "normalize": self.normalize,
        }

        import json
        with open(os.path.join(save_directory, "config.json"), "w") as f:
            json.dump(config, f, indent=2)

        print(f"Saved MizanEmbeddingModel → {save_directory}")

    # ------------------------------------------------------------
    # Load model
    # ------------------------------------------------------------
    @classmethod
    def from_pretrained(cls, folder, device="cpu"):
        import json

        config_path = os.path.join(folder, "config.json")
        with open(config_path, "r") as f:
            cfg = json.load(f)

        model = cls(
            backbone_name=cfg["backbone_name"],
            emb_dim=cfg["emb_dim"],
            pooling=cfg["pooling"],
            normalize=cfg["normalize"],
        )

        weights_path = os.path.join(folder, "pytorch_model.bin")
        state = torch.load(weights_path, map_location=device)
        model.load_state_dict(state)
        return model


# ============================================================
#  Wrapper with FIXED loading logic
# ============================================================
class MizanTextEncoderWrapper:
    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")

        # ------------------------------------------
        # Case 1: user gave a directory → Load custom
        # ------------------------------------------
        if os.path.isdir(backbone_name):
            print(f"Loading custom trained Mizan model from: {backbone_name}")

            # load model
            self.model = MizanEmbeddingModel.from_pretrained(
                backbone_name, device=self.device
            ).to(self.device)

            # load tokenizer of underlying backbone
            self.tokenizer = AutoTokenizer.from_pretrained(
                self.model.backbone_name
            )

        else:
            # ------------------------------------------
            # Case 2: user gave a huggingface model name
            # ------------------------------------------
            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:
        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:
        return self.encode([text])[0]
