# -*- coding: utf-8 -*-
"""
final lip-sync module (Optimized for HF Spaces)
"""
from huggingface_hub import hf_hub_download
import gradio as gr
import cv2
import torch
import torch.nn as nn
from torchvision import transforms
import librosa
import numpy as np
import numpy
# Naye PyTorch ko batana ke audio model ki weights safe hain
try:
    torch.serialization.add_safe_globals([np.core.multiarray.scalar])
except:
    pass

# --- 🚨 MASTER FIX FOR PYTORCH 2.6 SECURITY ---
_original_load = torch.load
def _patched_load(*args, **kwargs):
    kwargs['weights_only'] = False
    return _original_load(*args, **kwargs)
torch.load = _patched_load
# ----------------------------------------------

import os
import gc
import re
import requests
import base64
import random
from io import BytesIO
from PIL import Image
from sklearn.metrics.pairwise import cosine_similarity, euclidean_distances
from sklearn.preprocessing import normalize
import open_clip
import laion_clap
import whisper
import timm  

from transformers import (
    ViTForImageClassification, ViTImageProcessor,
    VideoMAEImageProcessor, VideoMAEForVideoClassification,
    AutoFeatureExtractor, AutoModelForAudioClassification, AutoModelForAudioFrameClassification,
    Wav2Vec2Processor, Wav2Vec2Model
)
from retinaface import RetinaFace
import torch.nn.functional as F
from groq import Groq

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

# API Keys (Securely fetched from HF Secrets)
OPENROUTER_API_KEY = os.environ.get("OPENROUTER_API_KEY")
GROQ_API_KEY = os.environ.get("GROQ_API_KEY")

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
client = Groq(api_key=GROQ_API_KEY)

# 👇 Change 1: Updated AST Model Path to your new Repo
AST_MODEL_PATH = "aneela-pervez/FAKE-AUDIO"
LIP_MODEL_PATH = hf_hub_download(repo_id="aneela-pervez/My-Deepfake-Models", filename="best_model (3).pth")

def clear_memory():
    gc.collect()
    if torch.cuda.is_available():
        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."
    ]
}

# CLAP SEMANTIC PROMPTS
CLAP_REAL_PROMPTS = [
    "A person speaking naturally",
    "Human voice with natural breathing and emotion",
    "Clear natural human conversation"
]
CLAP_FAKE_PROMPTS = [
    "Robotic AI generated voice",
    "Monotone synthetic speech without emotion",
    "Deepfake voice clone with metallic resonance"
]

# --- 3. AUDIO FORENSIC CATEGORIES ---
FORENSIC_CATEGORIES_AUDIO = {
    "AUTHENTIC": ["Natural human speech", "Human voice with natural breathing", "Consistent room ambiance"],
    "AI_SYNTHETIC": ["AI voice clone robotic smoothness", "Synthetic neural vocoder", "Deepfake metallic resonance"],
    "MANIPULATED": ["Manually spliced audio", "Inconsistent stitching artifacts", "Micro-cuts and clicks"],
    "DISGUISED": ["Artificial pitch shifting", "Muffled voice mask identity", "Time-stretched artifacts"]
}
ALL_AUDIO_LABELS = [item for sublist in FORENSIC_CATEGORIES_AUDIO.values() for item in sublist]

# --- 4. MODEL LOADING (ALL EXPERTS) ---
def load_all_models():
    print("Loading All Forensic Models...")

    try:
        vit_proc = ViTImageProcessor.from_pretrained("dima806/deepfake_vs_real_image_detection", token=os.environ.get("HF_TOKEN"))
        vit_model = ViTForImageClassification.from_pretrained("dima806/deepfake_vs_real_image_detection", token=os.environ.get("HF_TOKEN")).to(DEVICE).eval()

        vmae_path = "MCG-NJU/videomae-base-finetuned-kinetics"
        vmae_proc = VideoMAEImageProcessor.from_pretrained(vmae_path)
        vmae_model = VideoMAEForVideoClassification.from_pretrained(vmae_path).to(DEVICE).eval()

        clip_net, _, clip_tf = open_clip.create_model_and_transforms('ViT-B-32', pretrained='laion2b_s34b_b79k')
        clip_net.to(DEVICE).eval()

        w2v_proc = Wav2Vec2Processor.from_pretrained("facebook/wav2vec2-base-960h")
        w2v_model = Wav2Vec2Model.from_pretrained("facebook/wav2vec2-base-960h").to(DEVICE).eval()

        whisper_mod = whisper.load_model("base", device=DEVICE)

        id_ext = AutoFeatureExtractor.from_pretrained("facebook/wav2vec2-base-960h")
        id_mod = AutoModelForAudioFrameClassification.from_pretrained("facebook/wav2vec2-base-960h").to(DEVICE).eval()

        try:
            # 👇 Change 2: Removed local_files_only=True so it fetches from HF Hub
            ast_ext = AutoFeatureExtractor.from_pretrained(AST_MODEL_PATH)
            ast_mod = AutoModelForAudioClassification.from_pretrained(AST_MODEL_PATH).to(DEVICE).eval()
            print("✅ Loaded Custom AST Model")
        except:
            print("⚠️ Custom AST not found, using standard MIT/ast-finetuned-audioset")
            ast_ext = AutoFeatureExtractor.from_pretrained("MIT/ast-finetuned-audioset-10-10-0.4593")
            ast_mod = AutoModelForAudioClassification.from_pretrained("MIT/ast-finetuned-audioset-10-10-0.4593").to(DEVICE).eval()

        clap_mod = laion_clap.CLAP_Module(enable_fusion=False, amodel='HTSAT-tiny').to(DEVICE)
        clap_mod.load_ckpt(model_id=1)

        lip_mod = timm.create_model('legacy_xception', pretrained=False, num_classes=2)
        if os.path.exists(LIP_MODEL_PATH):
            try:
                cp = torch.load(LIP_MODEL_PATH, map_location=DEVICE, weights_only=False)
                state = cp.get('model_state_dict', cp)
                lip_mod.load_state_dict(state, strict=False)
                print("✅ Custom Lip-Sync Model Loaded")
            except Exception as e:
                print(f"⚠️ Custom Lip-Sync Model load failed: {e}. Using generic weights.")
        else:
            print("⚠️ Custom Lip-Sync Model not found at path, using generic weights.")
        lip_mod.to(DEVICE).eval()

        print("✅ All Models Loaded Successfully")
        return {
            "vit": (vit_model, vit_proc),
            "vmae": (vmae_model, vmae_proc),
            "clip": (clip_net, clip_tf),
            "w2v": (w2v_model, w2v_proc),
            "whisper": whisper_mod,
            "id": (id_mod, id_ext),
            "ast": (ast_mod, ast_ext),
            "clap": clap_mod,
            "lip": lip_mod
        }

    except Exception as e:
        print(f"❌ Critical Model Load Error: {e}")
        return None

models = load_all_models()

# --- 5. CLIP TRANSFORM ---
clip_transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711))
])

# --- 6. PROJECTION LAYER FOR LSE-NET ---
class LSEProjection(nn.Module):
    def __init__(self):
        super().__init__()
        self.proj = nn.Linear(400, 768)
    def forward(self, x):
        return self.proj(x)

lse_proj = LSEProjection().to(DEVICE)

# --- 7. HELPER FUNCTIONS ---
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

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))
    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()
    while len(frames) < 16: frames.append(frames[-1])
    return frames[:16]

# --- 8. AUDIO ANALYSIS MODULE ---
def analyze_audio_fuzzy(audio_path):
    try:
        y, sr = librosa.load(audio_path, sr=16000)
        y_48, _ = librosa.load(audio_path, sr=48000)
        
        # 👇 Change 3: Audio Normalization added here to fix Web Compression Noise
        y = librosa.util.normalize(y)
        y_48 = librosa.util.normalize(y_48)
        
        y = y / (np.max(np.abs(y)) + 1e-9)

        # A. Extract DNA (ID Match)
        id_mod, id_ext = models["id"]
        def get_dna(sig):
            inp = id_ext(sig, sampling_rate=16000, return_tensors="pt").to(DEVICE)
            with torch.no_grad():
                logits = id_mod(**inp).logits
                emb = torch.mean(logits, dim=1).cpu().numpy()
            return normalize(emb)

        segments = np.array_split(y, 3)
        seg_embs = [get_dna(s) for s in segments]
        cos_sim = cosine_similarity(seg_embs[0], seg_embs[-1])[0][0]

        # B. Feature Extraction (AST Texture)
        ast_mod, ast_ext = models["ast"]
        inp = ast_ext(y, sampling_rate=16000, return_tensors="pt").to(DEVICE)
        with torch.no_grad():
            ast_prob = torch.nn.functional.softmax(ast_mod(**inp).logits, dim=-1).cpu().numpy()[0][1]
            ast_score = ast_prob * 100

        # C. CLAP SEMANTIC PROMPT MATCHING
        clap_mod = models["clap"]
        audio_emb = torch.from_numpy(clap_mod.get_audio_embedding_from_data(x=[y_48])).to(DEVICE)
        real_emb = torch.from_numpy(clap_mod.get_text_embedding(CLAP_REAL_PROMPTS)).to(DEVICE)
        fake_emb = torch.from_numpy(clap_mod.get_text_embedding(CLAP_FAKE_PROMPTS)).to(DEVICE)

        sim_real = torch.nn.functional.cosine_similarity(audio_emb, real_emb).mean().item()
        sim_fake = torch.nn.functional.cosine_similarity(audio_emb, fake_emb).mean().item()

        sim_real_norm = max(0, sim_real + 0.1)
        sim_fake_norm = max(0, sim_fake + 0.1)
        clap_fake_score = (sim_fake_norm / (sim_real_norm + sim_fake_norm)) * 100

        # D. Edits / Micro-cuts
        num_edits = len(np.where(np.diff(librosa.onset.onset_strength(y=y, sr=16000)) > 8.0)[0])
        edit_penalty = max(0, (num_edits - 15) * 3.0)

        # E. Hybrid Final Anomaly Score
        final_anomaly_score = (clap_fake_score * 0.65) + (ast_score * 0.35) + edit_penalty
        final_anomaly_score = min(100.0, final_anomaly_score)

        metrics = f"CLAP Fake Match: {clap_fake_score:.1f}% | AST Texture: {ast_score:.1f}% | Cuts: {num_edits} | ID Match: {cos_sim:.2f}"

        if final_anomaly_score > 70.0:
            verdict = "FAKE"
            reason = "AI Voice Detected (CLAP Semantic Match)"
        else:
            verdict = "REAL"
            reason = "Authentic Voice (CLAP Verified)"

        return verdict, final_anomaly_score, reason, metrics

    except Exception as e:
        return "ERROR", 0, str(e), ""

# --- 9. INTEGRATED LIP SYNC MODULE ---
def analyze_lip_sync(frames, audio_path):
    try:
        speech, _ = librosa.load(audio_path, sr=16000)
        # 👇 Change 3 continued: Normalize lip sync audio too
        speech = librosa.util.normalize(speech)
        
        w2v_mod, w2v_proc = models["w2v"]

        a_inputs = w2v_proc(speech, return_tensors="pt", sampling_rate=16000).input_values.to(DEVICE)
        with torch.no_grad():
            audio_embeddings = w2v_mod(a_inputs).last_hidden_state.mean(dim=1)
            audio_embeddings = F.normalize(audio_embeddings, dim=-1) 

        vmae_mod, vmae_proc = models["vmae"]
        v_inputs = vmae_proc(list(frames), return_tensors="pt").to(DEVICE)

        with torch.no_grad():
            video_outputs = vmae_mod(**v_inputs).logits 
            video_embeddings = lse_proj(video_outputs)  
            video_embeddings = F.normalize(video_embeddings, dim=-1) 

        sync_similarity = F.cosine_similarity(audio_embeddings, video_embeddings).item()
        sync_error_score = (1 - max(0, sync_similarity)) * 100

        return sync_error_score

    except Exception as e:
        print(f"LSE-Net Lip-Sync Verification Failed: {e}")
        return 100.0 

# --- 10. MASTER PIPELINE ---
def master_pipeline(video_path):
    clear_memory()
    if not video_path: return None, "No Video", "Error"

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

    vmae_mod, vmae_proc = models["vmae"]
    v_in = vmae_proc(list(frames), return_tensors="pt").to(DEVICE)
    with torch.no_grad():
        v_conf = torch.softmax(vmae_mod(**v_in).logits, dim=1).max().item() * 100
        video_score = 100 - v_conf 

    vit_mod, vit_proc = models["vit"]
    face = extract_faces_retina(frames[8]) 
    v_inp = vit_proc(images=face, return_tensors="pt").to(DEVICE)
    with torch.no_grad():
        vit_prob = vit_mod(**v_inp).logits.softmax(dim=1)[0][0].item() * 100 

    audio_verdict, audio_score, audio_reason, audio_metrics = analyze_audio_fuzzy(video_path)
    sync_score = analyze_lip_sync(frames, video_path)

    VIDEO_FAKE_THRESH = 75
    AUDIO_FAKE_THRESH = 75
    SYNC_BAD_THRESH = 65

    is_video_fake = (video_score > VIDEO_FAKE_THRESH) or (vit_prob > 50)
    is_audio_fake = (audio_verdict != "REAL")

    if not is_video_fake and not is_audio_fake:
        if sync_score > SYNC_BAD_THRESH:
            final_case = "CASE 1: Real Video + Real Audio + Lip Sync Issue"
            verdict_short = "SYNC_MISMATCH"
        else:
            final_case = "CASE 5: Real Video + Real Audio + Authentic"
            verdict_short = "AUTHENTIC"
    elif not is_video_fake and is_audio_fake:
        final_case = "CASE 2: Real Video + Fake Audio"
        verdict_short = "AUDIO_FAKE"
    elif is_video_fake and not is_audio_fake:
        final_case = "CASE 3: Fake Video + Real Audio"
        verdict_short = "VIDEO_FAKE"
    else: 
        final_case = "CASE 4: Fake Video + Fake Audio"
        verdict_short = "FULL_DEEPFAKE"

    prompts = random.sample(FORENSIC_PROMPTS["DeepFake"], 3)
    prompt_text = (
        f"Role: Senior Forensic Judge. Analyze this case.\n"
        f"Data:\n"
        f"1. Video Anomaly: {video_score:.1f}% (High=Fake)\n"
        f"2. Audio Anomaly: {audio_score:.1f}% (Verdict: {audio_verdict})\n"
        f"3. Lip-Sync Error (Cross-Modal Expert): {sync_score:.1f}% (High=Tampered)\n"
        f"4. Detected Case: {final_case}\n\n"
        f"Task: Provide a professional forensic justification for determining '{final_case}'. "
        f"Mention specific artifacts derived from these scores. Keep it concise."
    )

    buffered = BytesIO()
    frames[8].save(buffered, format="JPEG")
    img_str = base64.b64encode(buffered.getvalue()).decode()

    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)
        claude_reasoning = response.json()['choices'][0]['message']['content']
    except:
        claude_reasoning = f"Automated Verdict: {final_case}. Metrics align with this classification."

    summary = (
        f"# 🛡️ FINAL VERDICT: {verdict_short}\n"
        f"## 📂 Classification: {final_case}\n\n"
        f"### 🤖 Forensic Reasoning (Claude):\n{claude_reasoning}\n\n"
        f"### 📊 Detailed Metrics:\n"
        f"- **Lip-Sync Anomaly (LSE-Net):** {sync_score:.1f}% (Threshold: {SYNC_BAD_THRESH})\n"
        f"- **Video Score:** {video_score:.1f}% (Threshold: {VIDEO_FAKE_THRESH})\n"
        f"- **Audio Score:** {audio_score:.1f}% (Threshold: {AUDIO_FAKE_THRESH})\n"
        f"- **Audio Details:** {audio_metrics}\n"
        f"- **Video Artifacts:** Checked {', '.join(prompts)}"
        f"Mention specific artifacts. STRICT LIMIT: Write exactly 1 short sentence (Maximum 15 to 20 words)."
    )

    return frames[8], summary, claude_reasoning

# 👇 Change 4: Removed UI launch from here since app.py will handle it