import gradio as gr
import torch
import mimetypes
from PIL import Image
import cv2
from torchvision.models import efficientnet_b0
from torchvision import transforms

# =========================
# Load Model
# =========================
def load_model():
    model = efficientnet_b0()
    model.classifier[1] = torch.nn.Linear(model.classifier[1].in_features, 2)
    model.load_state_dict(torch.load("models/best_model-v3.pt", map_location="cpu"))
    model.eval()
    return model

model = load_model()

# =========================
# Preprocessing
# =========================
preprocess = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(
        mean=[0.485, 0.456, 0.406],
        std=[0.229, 0.224, 0.225]
    )
])

# =========================
# Image Prediction
# =========================
def predict_image(path):
    img = Image.open(path).convert("RGB")
    tensor = preprocess(img).unsqueeze(0)

    with torch.no_grad():
        out = model(tensor)
        probs = torch.softmax(out, dim=1)[0]
        conf, pred = torch.max(probs, dim=0)

    label = "🟢 Real" if pred.item() == 0 else "🔴 Deepfake"
    return label, f"{conf.item()*100:.2f}%", img


# =========================
# Video Prediction (Every 10th Frame)
# =========================
def predict_video(path):
    cap = cv2.VideoCapture(path)

    frame_count = 0
    predictions = []
    preview_img = None

    while True:
        ret, frame = cap.read()
        if not ret:
            break

        # Process every 10th frame
        if frame_count % 10 == 0:
            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            img = Image.fromarray(frame_rgb)

            # Save first sampled frame for preview
            if preview_img is None:
                preview_img = img

            tensor = preprocess(img).unsqueeze(0)

            with torch.no_grad():
                out = model(tensor)
                probs = torch.softmax(out, dim=1)[0]
                predictions.append(probs)

        frame_count += 1

    cap.release()

    if len(predictions) == 0:
        return "❌ No valid frames found", "", None

    # Average all frame probabilities
    avg_probs = torch.stack(predictions).mean(dim=0)
    conf, pred = torch.max(avg_probs, dim=0)

    label = "🟢 Real (Multi-frame)" if pred.item() == 0 else "🔴 Deepfake (Multi-frame)"

    return label, f"{conf.item()*100:.2f}%", preview_img


# =========================
# Main Prediction Router
# =========================
def predict_file(file_obj):
    if file_obj is None:
        return "⚠️ No file selected", "", None

    path = file_obj.name
    mime, _ = mimetypes.guess_type(path)

    if mime and mime.startswith("image"):
        return predict_image(path)

    elif mime and mime.startswith("video"):
        return predict_video(path)

    else:
        return "Unsupported file type", "", None


# =========================
# Gradio UI
# =========================
with gr.Blocks(title="Deepfake Detector") as demo:
    gr.Markdown("## 🧠 Deepfake Detector\nUpload an image or video to analyze authenticity.")

    file_input = gr.File(
        label="Drop File Here",
        file_types=[".jpg", ".jpeg", ".png", ".mp4", ".mov"]
    )

    with gr.Row():
        prediction = gr.Textbox(label="Prediction", interactive=False)
        confidence = gr.Textbox(label="Confidence (%)", interactive=False)

    preview = gr.Image(label="Preview", interactive=False)

    file_input.change(
        fn=predict_file,
        inputs=file_input,
        outputs=[prediction, confidence, preview]
    )

demo.launch()