# -*- coding: utf-8 -*-
"""
Cleaned fake video detection module (Gradio UI & Notebook Magics removed)
"""

# Note: In pure Python files, install requirements via terminal.
# pip install transformers av

# Agar aap locally run kar rahi hain, toh drive mount ki zaroorat nahi hai:
# from google.colab import drive
# drive.mount('/content/drive')

import cv2
import torch
import torch.nn as nn
from torchvision import transforms
from PIL import Image
import open_clip
import requests
import base64
import random
from io import BytesIO
import numpy as np
import os
import timm
from transformers import ViTForImageClassification, ViTImageProcessor
from transformers import VideoMAEImageProcessor, VideoMAEForVideoClassification
from retinaface import RetinaFace
import torch.nn.functional as F
import gc
import re

# --- 1. SETUP & AUTH ---
print("🚀 Initializing VideoMAE Forensic System...")

# Mount Drive (Optional - Commented out for local use)
# if not os.path.exists('/content/drive'):
#     print("Mounting Google Drive...")
#     drive.mount('/content/drive')

# API Keys
# os.environ["HF_TOKEN"] = "YOUR_TOKEN_HERE"
OPENROUTER_API_KEY = "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"
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

def clear_memory():
    gc.collect()
    torch.cuda.empty_cache()

clear_memory()

# --- 2. FORENSIC PROMPTS ---
FORENSIC_PROMPTS = {
    "DeepFake": [
        "Check face resolution vs hair.", "Check chin blurring.", "Check eyebrow definition.",
        "Check blinking mechanics.", "Check skin tone mismatch.", "Check rectangular artifacts.",
        "Check wax-like skin.", "Check gaze alignment.", "Check teeth separation.",
        "Check double chin blur.", "Check glasses artifacts.", "Check face flickering.",
        "Check edge seams.", "Check flat lighting.", "Check texture repetition."
    ],
    "MotionAnomaly": [
        "Check for jittery head movements.", "Check if expressions freeze unnaturally.",
        "Check if eye blinking is too slow or fast.", "Check for morphing artifacts during movement.",
        "Check if the face slides over the background.", "Check mouth movement consistency."
    ]
}

# --- 3. LOAD MODELS (CORRECTED) ---
def load_models():
    print("Loading AI Models (FP16 Optimized)...")
    clear_memory()

    # A. ViT (Frame Quality Detector)
    try:
        vit_processor = ViTImageProcessor.from_pretrained("dima806/deepfake_vs_real_image_detection", token=os.environ["HF_TOKEN"])
        vit_model = ViTForImageClassification.from_pretrained("dima806/deepfake_vs_real_image_detection", token=os.environ["HF_TOKEN"])
        vit_model.to(DEVICE).half().eval()
        print("✅ ViT Loaded")
    except Exception as e:
        print(f"⚠️ ViT Load Failed: {e}")
        vit_model = None

    # B. VideoMAE (Motion Expert) - FIXED MODEL ID
    print("Loading VideoMAE (Stable V1)...")
    try:
        # CORRECT MODEL ID: MCG-NJU/videomae-base-finetuned-kinetics
        videomae_path = "MCG-NJU/videomae-base-finetuned-kinetics"
        videomae_processor = VideoMAEImageProcessor.from_pretrained(videomae_path)
        videomae_model = VideoMAEForVideoClassification.from_pretrained(videomae_path)
        videomae_model.to(DEVICE).half().eval()
        print("✅ VideoMAE Loaded (Motion Expert)")
    except Exception as e:
        print(f"⚠️ VideoMAE Failed: {e}")
        videomae_model = None
        videomae_processor = None

    # C. OpenCLIP (Semantic Engine)
    clip_net, _, clip_preprocess = open_clip.create_model_and_transforms('ViT-B-32', pretrained='laion2b_s34b_b79k')
    clip_net.to(DEVICE).eval()
    print("✅ OpenCLIP Loaded")

    return vit_model, vit_processor, videomae_model, videomae_processor, clip_net, clip_transform

# Dummy transform for CLIP
clip_transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711))
])

# Load models globally
vit_model, vit_processor, videomae_model, videomae_processor, clip_net, _ = load_models()

# --- 4. UTILS & PROCESSING ---
def extract_faces_retina(frame_img):
    frame_np = np.array(frame_img)
    try:
        faces = RetinaFace.detect_faces(frame_np)
    except: return frame_img

    if not faces or isinstance(faces, tuple): return frame_img

    max_area = 0
    best_face = frame_img
    for key in faces:
        identity = faces[key]
        facial_area = identity["facial_area"]
        x1, y1, x2, y2 = facial_area
        width = x2 - x1
        height = y2 - y1
        area = width * height
        if area > max_area:
            max_area = area
            margin = int(width * 0.2)
            x1 = max(0, x1 - margin)
            y1 = max(0, y1 - margin)
            x2 = min(frame_np.shape[1], x2 + margin)
            y2 = min(frame_np.shape[0], y2 + margin)
            best_face = frame_img.crop((x1, y1, x2, y2))
    return best_face

# --- 5. SEMANTIC ANALYSIS (COSINE) ---
def get_semantic_score(image, clip_model, device):
    real_descs = ["authentic high quality photo", "natural skin texture", "real human face"]
    fake_descs = ["deepfake artificial face", "robotic manipulated face", "blurry distorted deepfake"]

    real_tok = open_clip.get_tokenizer('ViT-B-32')(real_descs).to(device)
    fake_tok = open_clip.get_tokenizer('ViT-B-32')(fake_descs).to(device)
    img_tensor = clip_transform(image).unsqueeze(0).to(device)

    with torch.no_grad():
        img_emb = clip_model.encode_image(img_tensor)
        real_emb = clip_model.encode_text(real_tok)
        fake_emb = clip_model.encode_text(fake_tok)

        img_emb = F.normalize(img_emb, dim=-1)
        real_emb = F.normalize(real_emb, dim=-1)
        fake_emb = F.normalize(fake_emb, dim=-1)

        sim_real = (img_emb @ real_emb.T).mean().item()
        sim_fake = (img_emb @ fake_emb.T).mean().item()

    diff = sim_fake - sim_real
    return diff * 100

# --- 6. THE JUDGE (CLAUDE) ---
def judge_with_claude(image, vit_score, videomae_score, semantic_score):
    buffered = BytesIO()
    image.save(buffered, format="JPEG")
    img_str = base64.b64encode(buffered.getvalue()).decode()

    prompts = random.sample(FORENSIC_PROMPTS["DeepFake"], 3) + random.sample(FORENSIC_PROMPTS["MotionAnomaly"], 2)

    prompt_text = (
        f"Role: Forensic Video Judge. Determine if REAL or FAKE.\n"
        f"--- EVIDENCE PANEL ---\n"
        f"1. Semantic Analysis (Cosine): {semantic_score:.2f} (Positive=Fake)\n"
        f"2. VideoMAE (Motion Expert): {videomae_score:.1f}% Motion Anomaly Score\n"
        f"   *Note: High VideoMAE score means unnatural motion patterns.*\n"
        f"3. Visual Detector (ViT): {vit_score:.1f}% Fake\n\n"
        f"--- INSTRUCTIONS ---\n"
        f"1. Visually inspect artifacts: {', '.join(prompts)}.\n"
        f"2. Trust your eyes. If VideoMAE indicates high anomaly OR you see artifacts, declare FAKE.\n"
        f"3. Output REAL only if motion and pixels look natural.\n\n"
        f"--- FORMAT ---\n"
        f"FINAL_VERDICT: [REAL or FAKE]\n"
        f"FINAL_SCORE: [0-100]\n"
        f"EVIDENCE: [3 bullet points]"
    )

    payload = {
        "model": "anthropic/claude-3.5-sonnet",
        "messages": [{"role": "user", "content": [
            {"type": "text", "text": prompt_text},
            {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{img_str}"}}
        ]}]
    }

    try:
        response = requests.post("https://openrouter.ai/api/v1/chat/completions",
                               headers={"Authorization": f"Bearer {OPENROUTER_API_KEY}"},
                               json=payload)
        return response.json()['choices'][0]['message']['content']
    except Exception as e:
        return f"FINAL_VERDICT: ERROR\nFINAL_SCORE: 0\nEVIDENCE: {str(e)}"

# --- 7. MAIN PIPELINE ---
def extract_16_frames(video_path):
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened(): return []
    total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    # VideoMAE works best with 16 sampled frames
    indices = np.linspace(0, total-1, 16, dtype=int)
    frames = []
    for i in indices:
        cap.set(cv2.CAP_PROP_POS_FRAMES, i)
        ret, frame = cap.read()
        if ret: frames.append(Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)))
    cap.release()
    # Pad if video is too short
    while len(frames) < 16:
        frames.append(frames[-1])
    return frames[:16]

def parse_judge_output(text):
    verdict, score, evidence = "UNKNOWN", 0.0, text
    v_match = re.search(r"FINAL_VERDICT:\s*(REAL|FAKE)", text, re.IGNORECASE)
    if v_match: verdict = v_match.group(1).upper()
    s_match = re.search(r"FINAL_SCORE:\s*(\d+)", text)
    if s_match: score = float(s_match.group(1))
    e_match = re.search(r"EVIDENCE:\s*(.*)", text, re.DOTALL)
    if e_match: evidence = e_match.group(1).strip()
    return verdict, score, evidence

def pipeline(video):
    clear_memory()
    if video is None: return None, "No video", "", "Error"

    frames = extract_16_frames(video)
    if not frames: return None, "Extraction Failed", "", "Error"

    # 1. RUN VideoMAE (Batch Processing)
    videomae_score = 0
    if videomae_model:
        inputs = videomae_processor(list(frames), return_tensors="pt").to(DEVICE)
        inputs = {k: v.to(DEVICE).half() if v.dtype == torch.float else v for k, v in inputs.items()}
        with torch.no_grad():
            outputs = videomae_model(**inputs)
            logits = outputs.logits
            # VideoMAE is trained on Actions.
            # If the model is VERY confident about a standard action (e.g. "Talking"), it's likely Real.
            # If confidence is low/scattered, it implies Unnatural Motion (Fake).
            confidence = torch.softmax(logits, dim=1).max().item() * 100

            # We invert confidence: Low Confidence = High Anomaly Score
            # Deepfakes often lack the specific 'signatures' of natural Kinetics actions.
            videomae_score = 100 - confidence

    # 2. RUN ViT & Cosine (Frame by Frame)
    vit_scores = []
    semantic_scores = []
    processed_faces = []

    # Analyze fewer frames for visual checks to save time
    check_frames = [frames[0], frames[7], frames[15]]

    for frame in check_frames:
        face = extract_faces_retina(frame)
        processed_faces.append(face)

        with torch.no_grad():
            # ViT
            inputs = vit_processor(images=face, return_tensors="pt").to(DEVICE)
            inputs = {k: v.to(DEVICE).half() if v.dtype == torch.float else v for k, v in inputs.items()}
            vit_out = vit_model(**inputs)
            vit_scores.append(vit_out.logits.softmax(dim=1)[0][0].item() * 100) # 0=Fake

            # Semantic
            semantic_scores.append(get_semantic_score(face, clip_net, DEVICE))

    max_vit = max(vit_scores)
    avg_semantic = sum(semantic_scores) / len(semantic_scores)

    # Identify worst frame for Judge
    suspicious_idx = np.argmax(vit_scores)
    best_frame = processed_faces[suspicious_idx]

    # --- CLAUDE JUDGEMENT ---
    judge_response = judge_with_claude(best_frame, max_vit, videomae_score, avg_semantic)
    final_verdict, final_score, final_evidence = parse_judge_output(judge_response)

    if final_verdict == "UNKNOWN":
        final_verdict = "FAKE" if (max_vit > 50 or videomae_score > 60) else "REAL"
        final_evidence = judge_response

    emoji = "🚨" if final_verdict == "FAKE" else "✅"
    summary = (
        f"# {emoji} Final Verdict: {final_verdict}\n"
        f"**Judge Confidence:** {final_score:.1f}%\n"
        f"**(Source: Motion Anomaly + Visual Reasoning)**\n\n"
        f"**Model Breakdown:**\n"
        f"- 🏃 **VideoMAE (Motion Anomaly):** {videomae_score:.1f}% (High = Suspicious)\n"
        f"- 🧠 **Semantic (Cosine):** {avg_semantic:.2f}\n"
        f"- 👁️ **ViT Detector:** {max_vit:.1f}% Fake"
    )

    return best_frame, summary, final_evidence, final_verdict