import cv2
import numpy as np
import tensorflow as tf
import gradio as gr
import os
import tempfile
import traceback
from huggingface_hub import hf_hub_download

# =======================================
# Step 1: Load TFLite model from Hugging Face Hub
# =======================================
interpreter, input_details, output_details = None, None, None

try:
    model_path = hf_hub_download("Shakeel401/Deepfake-Detector", "deepfake_model_pure.tflite")

    interpreter = tf.lite.Interpreter(model_path=model_path)
    interpreter.allocate_tensors()
    input_details = interpreter.get_input_details()
    output_details = interpreter.get_output_details()
    print("✅ Loaded TFLite model from HF Hub")
except Exception as e:
    print("❌ Failed to load TFLite model:", e)

# =======================================
# Step 2: OpenCV face detector
# =======================================
face_cascade = cv2.CascadeClassifier(
    cv2.data.haarcascades + "haarcascade_frontalface_default.xml"
)

# =======================================
# Step 3: Preprocessing function
# =======================================
def preprocess_frames(frames, target_size=(299, 299)):
    clip_len = 10
    if len(frames) < clip_len:
        while len(frames) < clip_len:
            frames.append(frames[-1].copy())
    processed = []
    for frame in frames[:clip_len]:
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        faces = face_cascade.detectMultiScale(gray, scaleFactor=1.1, minNeighbors=5)
        if len(faces) > 0:
            x, y, w, h = faces[0]
            face = frame[y:y+h, x:x+w]
            face = cv2.resize(face, target_size)
        else:
            face = cv2.resize(frame, target_size)
        face = face.astype("float32") / 255.0
        processed.append(face)
    processed = np.array(processed, dtype=np.float32)
    processed = np.expand_dims(processed, 0)
    return processed

# =======================================
# Step 4: Inference wrapper (TFLite only)
# =======================================
def run_inference(processed):
    try:
        interpreter.set_tensor(input_details[0]['index'], processed)
        interpreter.invoke()
        return float(interpreter.get_tensor(output_details[0]['index'])[0][0])
    except Exception as e:
        print("❌ Inference failed:", e)
        return 0.5

# =======================================
# Step 5: Video prediction with overlay
# =======================================
MAX_VIDEO_SIZE_MB = 50

def predict_video_overlay(video_file):
    if interpreter is None:
        return "❌ No TFLite model available"

    try:
        file_size_mb = os.path.getsize(video_file)/(1024*1024)
        if file_size_mb > MAX_VIDEO_SIZE_MB:
            return f"❌ Video too large ({file_size_mb:.2f} MB)"
        
        cap = cv2.VideoCapture(video_file)
        frames_buffer, frame_count, results = [], 0, []
        temp_file = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)
        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
        fps = int(cap.get(cv2.CAP_PROP_FPS) or 25)
        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        out = cv2.VideoWriter(temp_file.name, fourcc, fps, (width, height))
        
        clip_len = 10
        frame_skip = 5
        
        while cap.isOpened():
            ret, frame = cap.read()
            if not ret:
                break
            if frame_count % frame_skip == 0:
                frames_buffer.append(frame)
                if len(frames_buffer) == clip_len:
                    processed = preprocess_frames(frames_buffer)
                    pred = run_inference(processed)
                    results.append(pred)
                    frames_buffer = []
            # overlay
            if results:
                avg_score = np.median(results)
                label = "Fake" if avg_score >= 0.5 else "Real"
                confidence = avg_score*100 if avg_score >= 0.5 else (1-avg_score)*100
                color = (0,0,255) if label=="Fake" else (0,255,0)
                text = f"{label} ({confidence:.2f}%)"
                cv2.putText(frame, text, (30,50), cv2.FONT_HERSHEY_SIMPLEX, 1, color, 2, cv2.LINE_AA)
            out.write(frame)
            frame_count += 1

        # final buffer
        if frames_buffer:
            processed = preprocess_frames(frames_buffer)
            pred = run_inference(processed)
            results.append(pred)

        cap.release()
        out.release()
        if not results:
            return "❌ No clips processed"
        return temp_file.name
    except Exception as e:
        traceback.print_exc()
        return f"❌ Error while processing video: {str(e)}"

# =======================================
# Step 6: Gradio App
# =======================================
title = "🎬 Deepfake Detection (TFLite from Hugging Face Hub)"
description = f"Upload MP4 (Max {MAX_VIDEO_SIZE_MB} MB). Detects Real/Fake and overlays predictions."

gr.Interface(
    fn=predict_video_overlay,
    inputs=gr.Video(label="Upload Video", sources=["upload"], format="mp4"),
    outputs=gr.Video(label="Video with Predictions"),
    title=title,
    description=description,
    flagging_mode="never"
).launch()