from contextlib import asynccontextmanager
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import logging

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

MODEL_ID = "vectara/hallucination_evaluation_model"

tokenizer = None
model = None


@asynccontextmanager
async def lifespan(app: FastAPI):
    global tokenizer, model
    logger.info(f"Loading model {MODEL_ID} ...")
    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
    model = AutoModelForSequenceClassification.from_pretrained(MODEL_ID)
    model.eval()
    logger.info("Model ready.")
    yield
    del model
    del tokenizer


app = FastAPI(title="HHEM API", version="1.0.0", lifespan=lifespan)


class ScoreRequest(BaseModel):
    source: str
    generated: str


class ScoreResponse(BaseModel):
    score: float
    label: str


@app.get("/health")
def health():
    return {"status": "ok", "model": MODEL_ID, "loaded": model is not None}


@app.post("/score", response_model=ScoreResponse)
def score(req: ScoreRequest):
    if model is None or tokenizer is None:
        raise HTTPException(status_code=503, detail="Model not loaded yet")

    inputs = tokenizer(
        req.source,
        req.generated,
        return_tensors="pt",
        truncation=True,
        max_length=512,
        padding=True,
    )

    with torch.no_grad():
        logits = model(**inputs).logits
        probs = torch.softmax(logits, dim=-1)

    # Resolve which class index maps to "Factually Consistent"
    id2label = model.config.id2label  # e.g. {0: "Hallucinated", 1: "Factually Consistent"}
    factual_idx = next(
        (idx for idx, name in id2label.items() if "consistent" in name.lower() or "factual" in name.lower()),
        1,  # fallback: assume class 1 is the positive (consistent) class
    )

    factual_score = probs[0][factual_idx].item()
    label = "grounded" if factual_score > 0.5 else "hallucinated"

    return ScoreResponse(score=round(factual_score, 4), label=label)