from fastapi import FastAPI, UploadFile, File, HTTPException
from fastapi.middleware.cors import CORSMiddleware
import shutil
import uuid
import os

from inference import run_inference, generate_heatmap

app = FastAPI()

@app.get("/")
def root():
    return {"status": "Deepfake API running"}

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

@app.post("/predict")
async def predict(video: UploadFile = File(...)):

    file_id = f"{uuid.uuid4()}.mp4"
    path = f"/tmp/{file_id}"

    try:

        with open(path, "wb") as buffer:
            shutil.copyfileobj(video.file, buffer)

        result = run_inference(path)

        return result

    except Exception as e:

        raise HTTPException(
            status_code=500,
            detail=str(e)
        )

    finally:

        if os.path.exists(path):
            os.remove(path)

@app.post("/heatmap")
async def heatmap(data: dict):

    try:

        frame_index = int(data.get("frame_index", -1))

        if frame_index < 0:
            raise HTTPException(
                status_code=400,
                detail="Invalid frame index"
            )

        heatmap, regions = generate_heatmap(frame_index)

        if heatmap is None:
            raise HTTPException(
                status_code=404,
                detail="Frame not found"
            )

        return {
            "heatmap": heatmap,
            "regions": regions
        }

    except Exception as e:

        raise HTTPException(
            status_code=500,
            detail=str(e)
        )