import os
import cv2
import torch
import shutil
import uuid
import yt_dlp
from typing import Optional
from fastapi import FastAPI, UploadFile, File, Form
from fastapi.middleware.cors import CORSMiddleware
from transformers import VideoMAEImageProcessor, VideoMAEForVideoClassification, pipeline
from facenet_pytorch import MTCNN
from PIL import Image

app = FastAPI()

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# ── Load models once at startup ───────────────────────────────────────────────
print("Loading MTCNN Face Detector...")
mtcnn = MTCNN(keep_all=False, select_largest=True, post_process=False)

print("Loading VideoMAE Temporal Deepfake Detector...")
model_name = "Ammar2k/videomae-base-finetuned-deepfake-subset"
processor  = VideoMAEImageProcessor.from_pretrained(model_name)
model      = VideoMAEForVideoClassification.from_pretrained(model_name)
model.eval()

print("Loading AI Image Detector (for fully synthetic AI-generated videos)...")
ai_image_detector = pipeline(
    "image-classification",
    model="Smogy/SMOGY-Ai-images-detector",
    device=-1  # CPU
)

os.makedirs("temp", exist_ok=True)

# ── Configuration ─────────────────────────────────────────────────────────────
SEQUENCE_LENGTH = 16   # VideoMAE requires exactly 16 frames

# IMPORTANT: This model is biased toward fake. Calibrated threshold after testing:
# Real videos score ~60-75%, so we raise the bar significantly.
FAKE_THRESHOLD  = 0.80  # Only call FAKE if model is 80%+ confident

# Threshold for the AI image detector (probability that a frame is AI-generated)
# Using a combo: flag if MAX single frame >= 0.55 OR average >= 0.30
AI_IMAGE_AVG_THRESHOLD = 0.30   # Flag if avg across all frames is >= 30%
AI_IMAGE_MAX_THRESHOLD = 0.55   # Flag if ANY single frame hits >= 55%
AI_FRAME_SAMPLES       = 8      # Number of frames to sample from the video

def smooth_box(current_box, last_box, alpha=0.5):
    """EMA smoothing to stabilise the face bounding box across frames."""
    if last_box is None:
        return current_box
    return [alpha * c + (1 - alpha) * p for c, p in zip(current_box, last_box)]

def crop_face(pil_img, box, padding=0.35):
    """Crop the face region with proportional padding."""
    w = box[2] - box[0]
    h = box[3] - box[1]
    pad_w = int(w * padding)
    pad_h = int(h * padding)
    x1 = max(0, int(box[0]) - pad_w)
    y1 = max(0, int(box[1]) - pad_h)
    x2 = min(pil_img.width,  int(box[2]) + pad_w)
    y2 = min(pil_img.height, int(box[3]) + pad_h)
    if x2 > x1 and y2 > y1:
        return pil_img.crop((x1, y1, x2, y2))
    return None

def extract_clip(video_path, start_frame):
    """
    Open a fresh cap, seek ONCE to start_frame, then read frames
    SEQUENTIALLY (no cap.set inside loop). This is reliable for all codecs.
    Collects SEQUENCE_LENGTH face crops using downscaled frame-by-frame tracking.
    """
    cap = cv2.VideoCapture(video_path)
    cap.set(cv2.CAP_PROP_POS_FRAMES, start_frame)  # Seek exactly ONCE

    faces    = []
    last_box = None
    attempts = 0

    while len(faces) < SEQUENCE_LENGTH and attempts < 60:
        ret, frame = cap.read()  # Sequential — no random seeking
        if not ret:
            break

        attempts += 1
        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        pil_img   = Image.fromarray(frame_rgb)

        # Downscale for faster MTCNN detection on every frame
        detect_img = pil_img.copy()
        current_box = None
        
        if detect_img.width > 640:
            ratio = 640.0 / detect_img.width
            detect_img = detect_img.resize((640, int(detect_img.height * ratio)))
            boxes, _ = mtcnn.detect(detect_img)
            if boxes is not None and len(boxes) > 0:
                current_box = [b / ratio for b in boxes[0].tolist()]
        else:
            boxes, _ = mtcnn.detect(pil_img)
            if boxes is not None and len(boxes) > 0:
                current_box = boxes[0].tolist()

        if current_box is not None:
            smoothed = smooth_box(current_box, last_box)
            last_box = smoothed
        elif last_box is not None:
            smoothed = last_box   # Hold last known position
        else:
            continue              # No face yet — keep reading

        # Use standard padding for centered faces
        crop = crop_face(pil_img, smoothed, padding=0.35)
        if crop is not None:
            faces.append(crop)

    cap.release()

    if not faces:
        return []

    while len(faces) < SEQUENCE_LENGTH:
        faces.append(faces[-1])  # Pad with last frame if clip was short

    return faces[:SEQUENCE_LENGTH]

def run_inference(faces):
    """
    Run VideoMAE on 16 face-crop frames.
    Returns the raw probability for the 'fake' class.
    """
    inputs = processor(list(faces), return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)

    probs = torch.nn.functional.softmax(outputs.logits[0], dim=-1)

    # Explicitly find the "fake" label index
    fake_prob = None
    for idx, label in model.config.id2label.items():
        if "fake" in label.lower():
            fake_prob = probs[idx].item()
            break

    # Fallback if label map is unexpected
    if fake_prob is None:
        predicted_idx = probs.argmax(-1).item()
        label = model.config.id2label[predicted_idx].lower()
        fake_prob = probs[predicted_idx].item() if "fake" in label else 1.0 - probs[predicted_idx].item()

    print(f"  id2label: {model.config.id2label}")
    print(f"  raw fake_prob: {fake_prob:.4f}")
    return fake_prob


def run_ai_image_check(video_path, total_frames):
    """
    Sample AI_FRAME_SAMPLES evenly-spaced frames from the video and run them
    through the AI image detector. Returns (is_ai_generated, avg_ai_score, triggered_frames).
    This catches fully synthetic videos (Gemini Veo, Sora, Runway, etc.) that
    VideoMAE misses because they have no face-swap artifacts.
    """
    step = max(1, total_frames // AI_FRAME_SAMPLES)
    frame_indices = [min(i * step, total_frames - 1) for i in range(AI_FRAME_SAMPLES)]

    ai_scores = []
    cap = cv2.VideoCapture(video_path)

    for idx in frame_indices:
        cap.set(cv2.CAP_PROP_POS_FRAMES, idx)
        ret, frame = cap.read()
        if not ret:
            continue
        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        pil_img = Image.fromarray(frame_rgb)

        results = ai_image_detector(pil_img)
        # Model labels vary; find the AI/Fake label score
        ai_score = 0.0
        for res in results:
            if any(kw in res['label'].lower() for kw in ['ai', 'fake', 'artificial', 'generated', 'synthetic']):
                ai_score = res['score']
                break
        ai_scores.append(ai_score)
        print(f"  Frame {idx}: AI image score = {ai_score:.4f}")

    cap.release()

    if not ai_scores:
        return False, 0.0, 0.0, 0

    avg_score = sum(ai_scores) / len(ai_scores)
    max_score = max(ai_scores)
    triggered = sum(1 for s in ai_scores if s >= AI_IMAGE_AVG_THRESHOLD)
    # Flag as AI-generated if avg is high OR any single frame was very strongly AI-detected
    is_ai = avg_score >= AI_IMAGE_AVG_THRESHOLD or max_score >= AI_IMAGE_MAX_THRESHOLD
    print(f"  AI Image Check — avg={avg_score:.4f}, max={max_score:.4f}, triggered={triggered}/{len(ai_scores)}, is_ai={is_ai}")
    return is_ai, avg_score, max_score, triggered

# ── API endpoint ──────────────────────────────────────────────────────────────
@app.post("/api/analyze")
async def analyze_video(
    file: Optional[UploadFile] = File(None),
    url: Optional[str] = Form(None)
):
    temp_path = None
    try:
        if file is not None and file.filename:
            print(f"Received file: {file.filename}")
            temp_path = os.path.join("temp", f"{uuid.uuid4()}_{file.filename}")
            with open(temp_path, "wb") as buf:
                shutil.copyfileobj(file.file, buf)
                
        elif url is not None and url.strip():
            print(f"Received URL: {url}")
            temp_id = str(uuid.uuid4())
            temp_path_template = os.path.join("temp", f"{temp_id}.%(ext)s")
            
            ydl_opts = {
                'format': 'best', # Simply download the best single file, avoiding ffmpeg merge requirements
                'outtmpl': temp_path_template,
                'noplaylist': True,
                'quiet': True,
                'max_filesize': 100 * 1024 * 1024  # Limit to 100MB
            }
            with yt_dlp.YoutubeDL(ydl_opts) as ydl:
                ydl.download([url])
                
            # Find the actual downloaded file since extension might vary
            for f in os.listdir("temp"):
                if temp_id in f:
                    temp_path = os.path.join("temp", f)
                    break
                    
            if not temp_path or not os.path.exists(temp_path):
                return {"error": "Failed to download the video from the provided URL."}

        else:
            return {"error": "Please provide either a video file or a valid URL."}

        cap          = cv2.VideoCapture(temp_path)
        video_fps    = cap.get(cv2.CAP_PROP_FPS) or 25.0
        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        cap.release()

        duration_s = total_frames / video_fps
        print(f"Video: {duration_s:.1f}s @ {video_fps:.1f}fps ({total_frames} frames)")

        if total_frames == 0:
            return {"isFake": False, "confidence": 0,
                    "explanation": "Could not read the video file.", "details": []}

        # ── Extract multiple clips from different parts of the video ───────────────
        segments = [
            max(0, int(total_frames * 0.2) - (SEQUENCE_LENGTH // 2)),
            max(0, int(total_frames * 0.5) - (SEQUENCE_LENGTH // 2)),
            max(0, int(total_frames * 0.8) - (SEQUENCE_LENGTH // 2))
        ]
        
        # Ensure unique starting frames if the video is very short
        segments = sorted(list(set(segments)))
        
        highest_fake_prob = 0.0
        successful_clips = 0

        for start_frame in segments:
            print(f"Extracting clip from frame {start_frame}...")
            faces = extract_clip(temp_path, start_frame)
            
            if faces:
                prob = run_inference(faces)
                highest_fake_prob = max(highest_fake_prob, prob)
                successful_clips += 1
                
                # If we already found very strong evidence of a fake, we can short-circuit to save time
                if highest_fake_prob > 0.95:
                    break

        # ── Run AI Image Detector on sampled frames ────────────────────────────
        # This catches fully synthetic AI-generated videos (Gemini, Sora, Runway, etc.)
        # that VideoMAE misses because they have no face-swap artifacts.
        print("Running AI Image Detector on sampled frames...")
        is_ai_generated, ai_avg_score, ai_max_score, ai_triggered = run_ai_image_check(temp_path, total_frames)

        # ── Combine both signals ───────────────────────────────────────────────
        # VideoMAE: catches face-swaps and traditional deepfakes
        # AI Image Detector: catches fully synthetic AI-generated content
        videomae_flagged = successful_clips > 0 and highest_fake_prob >= FAKE_THRESHOLD
        is_fake = videomae_flagged or is_ai_generated

        # Determine which method triggered and compute confidence
        if videomae_flagged and is_ai_generated:
            detection_method = "Dual-Model (Temporal + AI Image)"
            confidence = round(max(highest_fake_prob, ai_avg_score) * 100, 2)
            explanation = (
                "Both our Temporal VideoMAE and AI Image Detector flagged this video. "
                "It shows face-swap artifacts AND frame-level characteristics of AI-generated content."
            )
        elif is_ai_generated:
            detection_method = "AI Image Detector"
            confidence = round(max(ai_avg_score, ai_max_score) * 100, 2)
            explanation = (
                "Our AI Image Detector identified this video as fully synthetic — "
                f"frame-level analysis found strong AI-generation signatures "
                f"(peak score: {ai_max_score*100:.1f}%, avg: {ai_avg_score*100:.1f}%) "
                "consistent with tools like Gemini Veo, Sora, Runway, or similar generative AI systems."
            )
        elif videomae_flagged:
            detection_method = "VideoMAE Temporal Analysis"
            confidence = round(highest_fake_prob * 100, 2)
            explanation = (
                "Our Temporal AI detected strong evidence of facial manipulation — "
                "unnatural micro-expressions, blending artifacts, or temporal inconsistencies "
                "characteristic of deepfake face-swap synthesis."
            )
        else:
            # Neither triggered — real video
            detection_method = "Dual-Model"
            # Show highest confidence-of-real from both signals
            real_conf = max(
                (1.0 - highest_fake_prob) if successful_clips > 0 else 0.0,
                (1.0 - ai_avg_score)
            )
            confidence = round(real_conf * 100, 2)
            explanation = (
                "Our dual-model analysis found no significant manipulation. "
                "VideoMAE detected no temporal face-swap artifacts, and the AI Image Detector "
                "found no frame-level synthetic generation signatures."
            )

        print(f"FINAL — videomae={highest_fake_prob:.3f}, ai_avg={ai_avg_score:.3f}, ai_max={ai_max_score:.3f}, "
              f"method={detection_method}, isFake={is_fake}, confidence={confidence}%")

        return {
            "isFake": is_fake,
            "confidence": confidence,
            "explanation": explanation,
            "details": [
                {
                    "title": "VideoMAE Temporal Analysis",
                    "desc": f"Analyzed {successful_clips} clip(s) with a 3D VideoMAE Transformer. Peak score: {highest_fake_prob*100:.1f}%."
                },
                {
                    "title": "AI Image Frame Analysis",
                    "desc": f"Sampled {AI_FRAME_SAMPLES} frames for AI-generation signatures. Peak frame score: {ai_max_score*100:.1f}%, avg: {ai_avg_score*100:.1f}%. Detects Gemini Veo, Sora, Runway, etc."
                },
                {
                    "title": "Detection Method",
                    "desc": f"Result by: {detection_method}. Thresholds: VideoMAE ≥{FAKE_THRESHOLD*100:.0f}% | AI avg ≥{AI_IMAGE_AVG_THRESHOLD*100:.0f}% or max ≥{AI_IMAGE_MAX_THRESHOLD*100:.0f}%."
                }
            ]
        }

    except Exception as e:
        print(f"Error: {e}")
        import traceback; traceback.print_exc()
        return {"error": str(e)}

    finally:
        if temp_path and os.path.exists(temp_path):
            os.remove(temp_path)

@app.get("/")
def health_check():
    return {"status": "Calibrated VideoMAE Backend is running!"}