import gradio as gr
import torch
import cv2
import tempfile
import numpy as np
from PIL import Image
from transformers import AutoImageProcessor, AutoModelForImageClassification

# Load the deepfake detection model
processor = AutoImageProcessor.from_pretrained("Smogy/SMOGY-Ai-images-detector")
model = AutoModelForImageClassification.from_pretrained("Smogy/SMOGY-Ai-images-detector")

# Load face detection cascade
face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')

def detect_deepfake_image(image: Image.Image) -> str:
    """Detect deepfake in a single image"""
    inputs = processor(images=image, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)
    probs = torch.softmax(outputs.logits, dim=1)
    idx = probs.argmax().item()
    label = model.config.id2label[idx]
    conf = probs[0, idx].item()
    return f"The image is {label} with confidence {conf:.2f}"

def process_video(video_path: str) -> str:
    """Process video frame by frame and add detection annotations"""
    # Open input video
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        raise ValueError("Could not open video file")

    # Get video properties
    fps = cap.get(cv2.CAP_PROP_FPS)
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

    # Create temporary output file
    with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as temp_file:
        output_path = temp_file.name

    # Initialize video writer
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))

    # Process each frame
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break

        # Convert to grayscale for face detection
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        faces = face_cascade.detectMultiScale(gray, 1.1, 4)

        # Process each detected face
        for (x, y, w, h) in faces:
            # Extract face ROI
            face_img = frame[y:y+h, x:x+w]
            
            # Convert to PIL Image and process
            face_pil = Image.fromarray(cv2.cvtColor(face_img, cv2.COLOR_BGR2RGB))
            inputs = processor(images=face_pil, return_tensors="pt")
            
            with torch.no_grad():
                outputs = model(**inputs)
                
            probs = torch.softmax(outputs.logits, dim=1)
            idx = probs.argmax().item()
            label = model.config.id2label[idx]
            conf = probs[0, idx].item()

            # Draw bounding box and label
            color = (0, 255, 0) if label == 'real' else (0, 0, 255)  # BGR format
            cv2.rectangle(frame, (x, y), (x+w, y+h), color, 2)
            cv2.putText(frame, f"{label} {conf:.2f}", 
                        (x, y-10), cv2.FONT_HERSHEY_SIMPLEX, 
                        0.7, color, 2, cv2.LINE_AA)

        # Write processed frame
        out.write(frame)

    # Release resources
    cap.release()
    out.release()
    
    return output_path

with gr.Blocks() as demo:
    gr.Markdown("# Deepfake Detection Suite")
    
    with gr.Tab("Image Detection"):
        gr.Markdown("## Image Deepfake Detection")
        img_input = gr.Image(type="pil", label="Input Image")
        img_output = gr.Textbox(label="Detection Result")
        img_button = gr.Button("Analyze Image")
    
    with gr.Tab("Video Detection"):
        gr.Markdown("## Video Deepfake Detection")
        vid_input = gr.Video(label="Input Video")
        vid_output = gr.Video(label="Processed Video")
        vid_button = gr.Button("Analyze Video")
    
    img_button.click(fn=detect_deepfake_image, inputs=img_input, outputs=img_output)
    vid_button.click(fn=process_video, inputs=vid_input, outputs=vid_output)

if __name__ == "__main__":
    demo.launch()