"""
models.py  —  Hugging Face Hub edition
=======================================
This version loads your three fine-tuned BanglaBERT models from the
Hugging Face Hub (instead of Google Drive), which is what makes permanent
hosting on Hugging Face Spaces possible. app.py and the frontend are unchanged.

You only need to edit TWO things below:
    1. MODEL_IDS  -> your three model repositories on the Hub
    2. LABELS     -> the class names IN THE EXACT ORDER your model outputs them
"""

from __future__ import annotations
import os
import torch
import torch.nn.functional as F
from transformers import AutoTokenizer, AutoModelForSequenceClassification

# BanglaBERT's authors recommend normalizing Bangla text before inference.
try:
    from normalizer import normalize
except Exception:
    def normalize(text):
        return text

# Only needed if your model repositories are PRIVATE. On Hugging Face Spaces
# you set this by adding a secret named HF_TOKEN in the Space settings.
# For PUBLIC model repositories you can ignore this entirely.
HF_TOKEN = os.environ.get("HF_TOKEN")

# Fallback tokenizer if a repo somehow lacks its own tokenizer files.
BASE_BANGLABERT = "csebuetnlp/banglabert"

# ===========================================================================
# 1) YOUR THREE MODELS ON THE HUGGING FACE HUB.
#    Replace YOUR_USERNAME with your Hugging Face username. These must match
#    the repository names you created when uploading (see the guide, Part 2).
# ===========================================================================
MODEL_IDS = {
    "ai_detection": "Istihad/banglabert-ai-detection",
    "hate_speech":  "Istihad/banglabert-hate-speech",
    "sentiment":    "Istihad/banglabert-emotion",
}

# ===========================================================================
# 2) LABELS for each task, IN THE EXACT ORDER the model outputs them
#    (the class your model calls 0 goes first, then 1, and so on).
#    If a prediction ever looks wrong/flipped, fix the order here.
# ===========================================================================
LABELS = {
    "ai_detection": ["Human-written", "AI-generated"],
    "hate_speech":  ["Not hateful", "Hate / Offensive"],
    "sentiment":    ["Anger", "Disgust", "Fear", "Joy", "Sadness", "Surprise"],
}

# Human-readable names — also used to build the dropdown on the website.
TASKS = {
    "hate_speech": "Detect Bangla hate speech",
    "sentiment": "Detect emotion / mood",
    "ai_detection": "Detect AI-generated Bangla text",
}

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"


class HubModel:
    """Loads one fine-tuned BanglaBERT model from the Hub and runs predictions."""

    def __init__(self, task_key: str):
        self.task_key = task_key
        self.repo = MODEL_IDS[task_key]
        self.labels = LABELS[task_key]
        self.tokenizer = None
        self.model = None
        self._loaded = False

    def ensure_loaded(self):
        if self._loaded:
            return
        try:
            self.tokenizer = AutoTokenizer.from_pretrained(self.repo, token=HF_TOKEN)
        except Exception:
            self.tokenizer = AutoTokenizer.from_pretrained(BASE_BANGLABERT, token=HF_TOKEN)

        self.model = AutoModelForSequenceClassification.from_pretrained(
            self.repo, token=HF_TOKEN
        )
        self.model.to(DEVICE)
        self.model.eval()
        self._loaded = True

    def predict(self, text: str) -> dict:
        self.ensure_loaded()

        clean = normalize(text)
        inputs = self.tokenizer(
            clean, return_tensors="pt", truncation=True, max_length=512, padding=True
        ).to(DEVICE)

        with torch.no_grad():
            logits = self.model(**inputs).logits

        probs = F.softmax(logits, dim=-1)[0].cpu().tolist()
        best_idx = int(max(range(len(probs)), key=lambda i: probs[i]))

        details = {}
        for i, p in enumerate(probs):
            name = self.labels[i] if i < len(self.labels) else f"class_{i}"
            details[name] = round(p, 4)

        best_label = (
            self.labels[best_idx] if best_idx < len(self.labels) else f"class_{best_idx}"
        )
        return {"label": best_label, "score": float(probs[best_idx]), "details": details}


# One model object per task (weights download from the Hub on first use).
_REGISTRY = {key: HubModel(key) for key in MODEL_IDS}


def run_task(task_key: str, text: str) -> dict:
    """Entry point used by app.py."""
    if task_key not in _REGISTRY:
        raise ValueError(f"Unknown task: {task_key!r}")
    return _REGISTRY[task_key].predict(text)