from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from typing import Dict, List
from contextlib import asynccontextmanager
import torch
import logging
from datetime import datetime, timezone

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("sentiment-api")

MODEL_ID = "Whoawaisahmad/my-sentiment-analyzer"
ml_models: Dict[str, object] = {}

@asynccontextmanager
async def lifespan(app: FastAPI):
    logger.info(f"Loading model: {MODEL_ID}")
    try:
        tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
        model = AutoModelForSequenceClassification.from_pretrained(MODEL_ID)
        model.eval()
        ml_models["tokenizer"] = tokenizer
        ml_models["model"] = model
        logger.info("Model loaded successfully.")
    except Exception as e:
        logger.error(f"Failed to load model: {e}")
    yield
    ml_models.clear()

app = FastAPI(
    title="Sentiment Analysis API",
    description="A DistilBERT-based sentiment classifier.",
    version="1.0.0",
    lifespan=lifespan,
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

class TextRequest(BaseModel):
    text: str = Field(..., min_length=1, max_length=5000)

class BatchRequest(BaseModel):
    texts: List[str] = Field(..., min_length=1, max_length=50)

def predict(text):
    tokenizer = ml_models["tokenizer"]
    model = ml_models["model"]
    inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=128)
    with torch.no_grad():
        outputs = model(**inputs)
    probs = torch.softmax(outputs.logits, dim=-1)
    label = "POSITIVE" if probs[0][1] > probs[0][0] else "NEGATIVE"
    confidence = round(probs[0].max().item(), 4)
    other = round(1 - confidence, 4)
    return label, confidence, other

@app.get("/")
async def root():
    return {"message": "Sentiment Analysis API", "docs": "/docs", "health": "/health"}

@app.get("/health")
async def health_check():
    return {
        "status": "healthy" if ml_models.get("model") is not None else "degraded",
        "model_loaded": ml_models.get("model") is not None,
        "model_id": MODEL_ID,
        "timestamp": datetime.now(timezone.utc).isoformat(),
    }

@app.post("/predict")
async def predict_sentiment(request: TextRequest):
    if ml_models.get("model") is None:
        raise HTTPException(status_code=503, detail="Model not loaded.")
    try:
        label, confidence, other = predict(request.text)
        other_label = "NEGATIVE" if label == "POSITIVE" else "POSITIVE"
        return {
            "text": request.text[:200],
            "label": label,
            "confidence": confidence,
            "probabilities": {label: confidence, other_label: other},
            "timestamp": datetime.now(timezone.utc).isoformat(),
        }
    except Exception as e:
        logger.error(f"Prediction error: {e}")
        raise HTTPException(status_code=500, detail="Prediction failed.")

@app.post("/predict/batch")
async def predict_batch(request: BatchRequest):
    if ml_models.get("model") is None:
        raise HTTPException(status_code=503, detail="Model not loaded.")
    try:
        results = []
        for text in request.texts:
            label, confidence, _ = predict(text)
            results.append({"text": text[:100], "label": label, "confidence": confidence})
        return {"results": results, "total": len(results)}
    except Exception as e:
        logger.error(f"Batch error: {e}")
        raise HTTPException(status_code=500, detail="Batch prediction failed.")