from fastapi import FastAPI, Request
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
import uvicorn

app = FastAPI(title="AI Detector API")

# Load the model once at startup
MODEL_NAME = "roberta-base-openai-detector"
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME)
model.eval()

def get_ai_probability(text: str) -> float:
    """Return AI probability (0–100%) for the given text."""
    inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=512)
    with torch.no_grad():
        logits = model(**inputs).logits
        probs = torch.softmax(logits, dim=1)
    ai_score = probs[0][1].item() * 100
    return round(ai_score, 2)

@app.post("/analyze")
async def analyze_text(request: Request):
    """
    Example body:
    {
      "text": "Your text here"
    }
    """
    data = await request.json()
    text = data.get("text", "").strip()
    if not text:
        return {"error": "No text provided"}

    paragraphs = [p.strip() for p in text.split("\n") if p.strip()]
    results = []

    for i, para in enumerate(paragraphs, start=1):
        ai_score = get_ai_probability(para)
        results.append({
            "paragraph": i,
            "ai_score": ai_score,
            "human_score": round(100 - ai_score, 2),
            "content_preview": para[:200] + ("..." if len(para) > 200 else "")
        })

    overall = sum([r["ai_score"] for r in results]) / len(results)
    return {
        "overall_ai_score": round(overall, 2),
        "overall_human_score": round(100 - overall, 2),
        "paragraphs": results
    }

@app.get("/")
async def root():
    return {"message": "AI Detector API is running. Use POST /analyze"}

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