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

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

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

@asynccontextmanager
async def lifespan(app: FastAPI):
    logger.info(f"Loading model: {MODEL_ID}")
    try:
        ml_models["classifier"] = pipeline(
            "sentiment-analysis", model=MODEL_ID,
            device=0 if torch.cuda.is_available() else -1,
        )
        logger.info("Model loaded.")
    except Exception as e:
        logger.error(f"Failed: {e}")
        ml_models["classifier"] = None
    yield
    ml_models.clear()

app = FastAPI(title="Sentiment Analysis API", version="1.0.0", lifespan=lifespan)
app.add_middleware(CORSMiddleware, allow_origins=["*"], 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)

class SentimentResponse(BaseModel):
    text: str
    label: str
    confidence: float
    probabilities: Dict[str, float]
    timestamp: str

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

@app.get("/health")
async def health():
    loaded = ml_models.get("classifier") is not None
    return {"status": "healthy" if loaded else "degraded", "model_loaded": loaded}

@app.post("/predict", response_model=SentimentResponse)
async def predict(request: TextRequest):
    clf = ml_models.get("classifier")
    if clf is None:
        raise HTTPException(503, "Model not loaded.")
    result = clf(request.text)[0]
    other = "LABEL_0" if result["label"] == "LABEL_1" else "LABEL_1"
    return SentimentResponse(
        text=request.text[:200],
        label=result["label"],
        confidence=round(result["score"], 4),
        probabilities={result["label"]: round(result["score"], 4), other: round(1 - result["score"], 4)},
        timestamp=datetime.now(timezone.utc).isoformat(),
    )

@app.post("/predict/batch")
async def predict_batch(request: BatchRequest):
    clf = ml_models.get("classifier")
    if clf is None:
        raise HTTPException(503, "Model not loaded.")
    results = clf(request.texts)
    return {
        "results": [{"text": t[:100], "label": r["label"], "confidence": round(r["score"], 4)}
                    for t, r in zip(request.texts, results)],
        "total": len(results),
    }

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)