import os
import cv2
import math
import torch
import gradio as gr
import numpy as np
from transformers import AutoImageProcessor, AutoModelForVideoClassification

MODEL_REPO = "Vansh180/VideoMae-ffc23-deepfake-detector"
HF_TOKEN = os.getenv("HF_TOKEN", None)

device = "cuda" if torch.cuda.is_available() else "cpu"

processor = AutoImageProcessor.from_pretrained(MODEL_REPO, token=HF_TOKEN)
model = AutoModelForVideoClassification.from_pretrained(MODEL_REPO, token=HF_TOKEN)
model.to(device)
model.eval()

id2label = model.config.id2label

def sample_video_frames(video_path, num_frames=16):
    cap = cv2.VideoCapture(video_path)
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))

    if total_frames <= 0:
        cap.release()
        raise ValueError("Could not read video or empty video.")

    indices = np.linspace(0, total_frames - 1, num_frames, dtype=int)
    frames = []

    current_idx = 0
    target_set = set(indices.tolist())

    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break
        if current_idx in target_set:
            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            frames.append(frame)
        current_idx += 1

    cap.release()

    if len(frames) == 0:
        raise ValueError("No frames extracted from video.")

    while len(frames) < num_frames:
        frames.append(frames[-1])

    return frames[:num_frames]

def predict_video(video_path):
    if video_path is None:
        return {"error": "No video uploaded"}

    try:
        frames = sample_video_frames(video_path, num_frames=16)

        inputs = processor(frames, return_tensors="pt")
        inputs = {k: v.to(device) for k, v in inputs.items()}

        with torch.no_grad():
            outputs = model(**inputs)
            probs = torch.softmax(outputs.logits, dim=-1)[0]

        pred_idx = int(torch.argmax(probs).item())
        pred_label = id2label[pred_idx]
        confidence = float(probs[pred_idx].item())

        scores = {
            id2label[i]: float(probs[i].item())
            for i in range(len(probs))
        }

        return {
            "prediction": pred_label,
            "confidence": confidence,
            "scores": scores
        }

    except Exception as e:
        return {"error": str(e)}

demo = gr.Interface(
    fn=predict_video,
    inputs=gr.Video(label="Upload Video"),
    outputs=gr.JSON(label="Deepfake Detection Result"),
    title="Deepfake Video Detection",
    description="Upload a video to classify it as real or fake."
)

from fastapi import UploadFile, File
from fastapi.responses import JSONResponse
import tempfile
import os

@demo.app.post("/api/predict")
async def api_predict(file: UploadFile = File(...)):
    """Upload + predict in one call — for ShieldSense app."""
    if not file.filename or not any(file.filename.lower().endswith(ext) for ext in ('.mp4', '.mov', '.avi', '.webm', '.mkv', '.jpg', '.jpeg', '.png')):
        return JSONResponse({"error": "Invalid file type. Use video or image."}, status_code=400)
    try:
        with tempfile.NamedTemporaryFile(delete=False, suffix=os.path.splitext(file.filename)[1]) as f:
            f.write(await file.read())
            tmp_path = f.name
        try:
            result = predict_video(tmp_path)
            return JSONResponse(result)
        finally:
            os.unlink(tmp_path)
    except Exception as e:
        return JSONResponse({"error": str(e)}, status_code=500)

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)