from fastapi import FastAPI, UploadFile, File, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from PIL import Image
import io
from transformers import pipeline

app = FastAPI()

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

# Load AI Detection Model
# We use a pre-trained model from Hugging Face
# 'Organika/sdxl-detector' is specialized for Stable Diffusion detection
print("Loading AI Model...")
classifier = pipeline("image-classification", model="Organika/sdxl-detector")
print("Model Loaded!")

@app.post("/analyze")
async def analyze_media(file: UploadFile = File(...)):
    try:
        contents = await file.read()
        image = Image.open(io.BytesIO(contents))
        
        # Run Inference
        results = classifier(image)
        # results is a list like [{'label': 'artificial', 'score': 0.99}, {'label': 'human', 'score': 0.01}]
        
        # Find the 'artificial' or 'AI' score
        ai_score = 0.0
        for r in results:
            label = r['label'].lower()
            if 'artificial' in label or 'ai' in label:
                ai_score = r['score']
                break
            if 'human' in label or 'real' in label:
                # If we found human score, AI score is 1 - human
                ai_score = 1.0 - r['score']
        
        is_ai = ai_score > 0.5
        
        return {
            "filename": file.filename,
            "is_ai": is_ai,
            "confidence": round(ai_score * 100, 2), # Return 0-100
            "details": results
        }
    except Exception as e:
        print(f"Error: {e}")
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/")
def read_root():
    return {"status": "AI Detector Neural Network is Running"}