from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from fastapi.middleware.cors import CORSMiddleware
from transformers import pipeline
from PIL import Image
import base64
import io
import requests

app = FastAPI(title="STOA Deepfake Detector API")

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

# --- MODEL LOADING ---
print("Booting Security Node. Loading Deepfake ViT model into memory...")
pipe = pipeline("image-classification", model="prithivMLmods/deepfake-detector-model-v1")
print("Agent Ready!")

# --- REQUEST SCHEMA ---
class PredictRequest(BaseModel):
    image: str | None = None
    image_url: str | None = None

# --- ENDPOINTS ---
@app.get("/health")
def health_check():
    return {"status": "ok"}

@app.post("/predict")
def predict(req: PredictRequest):
    try:
        img = None
        
        if req.image_url:
            # The "Super-Human" header pack
            headers = {
                "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36",
                "Accept": "image/avif,image/webp,image/apng,image/*,*/*;q=0.8",
                "Referer": "https://google.com" # This tricks the site into thinking you came from a search
            }
            
            response = requests.get(req.image_url, stream=True, headers=headers, timeout=10)
            
            if response.status_code != 200:
                raise Exception(f"The image site blocked us with error: {response.status_code}. Try a different URL or Base64 upload.")
                
            img = Image.open(response.raw).convert("RGB")
        # 2. Handle Base64 Input
        elif req.image:
            b64_data = req.image
            if "," in b64_data:
                b64_data = b64_data.split(",")[1]
            image_bytes = base64.b64decode(b64_data)
            img = Image.open(io.BytesIO(image_bytes)).convert("RGB")
            
        # 3. Handle Empty Request
        else:
            raise HTTPException(status_code=400, detail="Must provide 'image' (base64) or 'image_url'.")

        # 4. Execute AI Math
        # 4. Execute AI Math
        results = pipe(img)
        
        # 5. Sensitivity Logic (The Hackathon Fix)
        # We manually check the fake score. If it's > 0.01, we flag it.
        fake_score = 0.0
        real_score = 0.0
        
        for res in results:
            if "FAKE" in res['label'].upper():
                fake_score += res['score']
            else:
                real_score += res['score']

        # THE FIX: Even if the AI says 90% Real, if it's 10% Fake, we flag it.
        # This is a common strategy for detecting high-quality modern deepfakes.
        if fake_score > 0.10:  # Adjust this (0.01 to 0.10) to find the sweet spot
            top_pred = "DEEPFAKE"
        else:
            top_pred = "REAL"

        return {
            "prediction": top_pred,
            "confidence": round(max(fake_score, real_score), 4),
            "scores": {
                "REAL": round(real_score, 4),
                "DEEPFAKE": round(fake_score, 4)
            }
        }
    except Exception as e:
        raise HTTPException(status_code=400, detail=f"Failed to process face: {str(e)}")