import os
import cv2
import torch
import zipfile
import librosa
import time
import subprocess
import tempfile
import numpy as np
import tensorflow as tf
from transformers import AutoFeatureExtractor, AutoModelForAudioClassification

try:
    import noisereduce as nr
    NOISEREDUCE_AVAILABLE = True
except ImportError:
    NOISEREDUCE_AVAILABLE = False

# Set random seed for reproducibility.
tf.random.set_seed(42)

# Extract EfficientNet model if not already extracted
if not os.path.exists("efficientnet-b0"):
    local_zip = "./efficientnet-b0.zip"
    if os.path.exists(local_zip):
        zip_ref = zipfile.ZipFile(local_zip, 'r')
        zip_ref.extractall()
        zip_ref.close()
        print("EfficientNet model extracted successfully!")

# Load EfficientNet model (image/video)
efficientnet_model = tf.keras.layers.TFSMLayer(
    "efficientnet-b0/",
    call_endpoint="serving_default"
)

# ─────────────────────────────────────────────────────────────────────────────
# Audio Ensemble: 3 models vote — majority wins (for uploaded files only)
# ─────────────────────────────────────────────────────────────────────────────
AUDIO_MODELS = [
    "MelodyMachine/Deepfake-audio-detection-V2",
    "MelodyMachine/Deepfake-audio-detection",
    "Gustking/wav2vec2-large-xlsr-deepfake-audio-classification",
]
AUDIO_SAMPLE_RATE = 16000

# ─── Model Thresholds ────────────────────────────────────────────────────────
REAL_THRESHOLD = 0.55
FAKE_THRESHOLD = 0.70

print("Loading audio ensemble models ...")
ensemble = []
for model_id in AUDIO_MODELS:
    print(f"  Loading {model_id} ...")
    try:
        fe = AutoFeatureExtractor.from_pretrained(model_id)
        m = AutoModelForAudioClassification.from_pretrained(model_id)
        m.eval()
        ensemble.append({"id": model_id, "extractor": fe, "model": m})
        print(f"  ✅ Loaded: {model_id} | labels: {m.config.id2label}")
    except Exception as e:
        print(f"  ⚠️ Skipped {model_id}: {e}")

print(f"Ensemble ready with {len(ensemble)} models.")


# ─────────────────────────────────────────────────────────────────────────────
# ACOUSTIC FEATURE ANALYZER
#
# Why do we need this?
#   All Wav2Vec2 models are binary (real/fake) — they cannot distinguish
#   AI synthesized audio from real because TTS doesn't match their "fake"
#   training patterns (replay attacks, splicing). They score TTS as "real".
#
# How does it work?
#   Real human voices have natural imperfections:
#     - Energy fluctuates (breathing, stress, pauses)
#     - Pitch varies naturally (prosody, emotion)
#     - Background noise / room acoustics present
#     - Zero crossing rate is irregular
#
#   AI synthesized voices are "too perfect":
#     - Energy is unnaturally consistent (flat amplitude envelope)
#     - Pitch follows mathematical patterns, low variance
#     - Very high SNR — almost no background noise
#     - Spectral flatness is high (energy distributed evenly)
#
# Decision:
#   acoustic_score = weighted combination of 4 features
#   score > AI_SYNTH_THRESHOLD → flag as AI Synthesized
#   This overrides a "real" vote from the model ensemble
# ─────────────────────────────────────────────────────────────────────────────

# Tune these thresholds based on testing:
# Higher = less sensitive (more audio passes as Real)
# Lower  = more sensitive (more audio flagged as AI Synthesized)
AI_SYNTH_THRESHOLD = 0.60  # overall acoustic score above this → AI Synthesized


def analyze_acoustic_features(x: np.ndarray, sr: int) -> dict:
    """
    Analyze audio for signs of AI synthesis by measuring naturalness.

    Returns a dict with individual feature scores (0=natural, 1=synthetic)
    and an overall ai_synth_score.
    """
    # ── Feature 1: Energy variance ────────────────────────────────────────────
    # Real voices: high energy variance (loud/quiet moments, breaths)
    # AI voices: low energy variance (flat, consistent loudness)
    frame_length = 1024
    hop_length = 256
    rms = librosa.feature.rms(y=x, frame_length=frame_length, hop_length=hop_length)[0]
    rms_variance = np.var(rms)
    rms_mean = np.mean(rms) + 1e-8

    # Normalize by mean energy — low coefficient of variation = synthetic
    rms_cv = np.sqrt(rms_variance) / rms_mean  # coefficient of variation
    # Typical real voice: cv > 0.5 | AI voice: cv < 0.3
    energy_synth_score = max(0.0, min(1.0, 1.0 - (rms_cv / 0.5)))
    print(f"[Acoustic] Energy CoV={rms_cv:.4f} → synth_score={energy_synth_score:.4f}")

    # ── Feature 2: Spectral flatness ─────────────────────────────────────────
    # Real voices: low spectral flatness (energy concentrated in harmonics)
    # AI voices: higher spectral flatness (more evenly distributed energy)
    spec_flatness = librosa.feature.spectral_flatness(y=x, hop_length=hop_length)[0]
    mean_flatness = np.mean(spec_flatness)
    # Typical real voice: < 0.05 | AI voice: > 0.08
    flatness_synth_score = max(0.0, min(1.0, mean_flatness / 0.1))
    print(f"[Acoustic] Spectral flatness={mean_flatness:.5f} → synth_score={flatness_synth_score:.4f}")

    # ── Feature 3: Pitch variance ─────────────────────────────────────────────
    # Real voices: pitch varies naturally with speech rhythm
    # AI voices: pitch follows smooth mathematical curves, lower variance
    try:
        f0 = librosa.yin(x, fmin=50, fmax=500, sr=sr, hop_length=hop_length)
        voiced = f0[f0 > 0]
        if len(voiced) > 10:
            pitch_variance = np.std(voiced) / (np.mean(voiced) + 1e-8)
            # Typical real voice: std/mean > 0.15 | AI voice: < 0.08
            pitch_synth_score = max(0.0, min(1.0, 1.0 - (pitch_variance / 0.15)))
        else:
            pitch_synth_score = 0.5  # not enough voiced frames to judge
    except Exception:
        pitch_synth_score = 0.5
    print(f"[Acoustic] Pitch variance score={pitch_synth_score:.4f}")

    # ── Feature 4: Zero Crossing Rate variance ────────────────────────────────
    # Real voices: ZCR fluctuates with consonants/vowels/pauses
    # AI voices: ZCR is more regular
    zcr = librosa.feature.zero_crossing_rate(x, hop_length=hop_length)[0]
    zcr_variance = np.var(zcr)
    zcr_mean = np.mean(zcr) + 1e-8
    zcr_cv = np.sqrt(zcr_variance) / zcr_mean
    # Typical real voice: cv > 0.5 | AI voice: cv < 0.3
    zcr_synth_score = max(0.0, min(1.0, 1.0 - (zcr_cv / 0.5)))
    print(f"[Acoustic] ZCR CoV={zcr_cv:.4f} → synth_score={zcr_synth_score:.4f}")

    # ── Weighted overall score ────────────────────────────────────────────────
    # Energy and pitch variance are most reliable indicators — weight them more
    ai_synth_score = (
        energy_synth_score  * 0.35 +
        flatness_synth_score * 0.20 +
        pitch_synth_score   * 0.30 +
        zcr_synth_score     * 0.15
    )
    print(f"[Acoustic] Overall AI synth score={ai_synth_score:.4f} (threshold={AI_SYNTH_THRESHOLD})")

    return {
        "energy_synth_score":   energy_synth_score,
        "flatness_synth_score": flatness_synth_score,
        "pitch_synth_score":    pitch_synth_score,
        "zcr_synth_score":      zcr_synth_score,
        "ai_synth_score":       ai_synth_score,
        "is_ai_synthesized":    ai_synth_score > AI_SYNTH_THRESHOLD,
    }


def convert_to_mp4(input_path):
    ext = os.path.splitext(input_path)[-1].lower()
    if ext == ".mp4":
        cap = cv2.VideoCapture(input_path)
        ok = cap.isOpened()
        cap.release()
        if ok:
            return input_path, False

    tmp = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)
    tmp.close()
    output_path = tmp.name

    cmd = [
        "ffmpeg", "-y", "-i", input_path,
        "-c:v", "libx264", "-preset", "fast",
        "-crf", "23", "-c:a", "aac", output_path
    ]
    result = subprocess.run(cmd, capture_output=True)
    if result.returncode != 0:
        os.unlink(output_path)
        raise RuntimeError(f"ffmpeg conversion failed:\n{result.stderr.decode()}")
    return output_path, True


class DetectionPipeline:
    def __init__(self, n_frames=None, batch_size=60, resize=None, input_modality='video'):
        self.n_frames = n_frames
        self.batch_size = batch_size
        self.resize = resize
        self.input_modality = input_modality

    def __call__(self, filename):
        if self.input_modality == 'video':
            print('Input modality is video.')
            converted_path, is_temp = convert_to_mp4(filename)

            try:
                v_cap = cv2.VideoCapture(converted_path)
                if not v_cap.isOpened():
                    raise RuntimeError(f"OpenCV could not open video: {converted_path}")

                v_len = int(v_cap.get(cv2.CAP_PROP_FRAME_COUNT))
                if v_len == 0:
                    raise RuntimeError("Video has 0 frames after conversion.")

                sample = (
                    np.arange(0, v_len) if self.n_frames is None
                    else np.linspace(0, v_len - 1, self.n_frames).astype(int)
                )

                faces, frames = [], []
                for j in range(v_len):
                    v_cap.grab()
                    if j in sample:
                        success, frame = v_cap.retrieve()
                        if not success:
                            continue
                        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                        if self.resize is not None:
                            frame = frame.resize([int(d * self.resize) for d in frame.size])
                        frames.append(frame)
                        if len(frames) % self.batch_size == 0 or j == sample[-1]:
                            faces.append(cv2.resize(frame, (224, 224)))
                v_cap.release()
            finally:
                if is_temp and os.path.exists(converted_path):
                    os.unlink(converted_path)

            if len(faces) == 0:
                raise RuntimeError("No frames could be extracted from the video.")
            return faces

        elif self.input_modality == 'image':
            image = cv2.cvtColor(filename, cv2.COLOR_BGR2RGB)
            return cv2.resize(image, (224, 224))

        else:
            raise ValueError(f"Invalid input modality: {self.input_modality}")


detection_video_pipeline = DetectionPipeline(n_frames=5, batch_size=1, input_modality='video')
detection_image_pipeline = DetectionPipeline(batch_size=1, input_modality='image')


def deepfakes_video_predict(input_video):
    faces = detection_video_pipeline(input_video)
    real_res, fake_res = [], []

    for face in faces:
        face2 = face / 255
        pred = efficientnet_model(np.expand_dims(face2, axis=0))
        pred = list(pred.values())[0].numpy()[0]
        real_res.append(pred[0])
        fake_res.append(pred[1])

    real_mean = np.mean(real_res)
    fake_mean = np.mean(fake_res)
    print(f"[Video] Real={real_mean:.4f} | Fake={fake_mean:.4f}")

    if real_mean >= 0.5:
        return "✅ The video is REAL."
    else:
        return "🚨 The video is FAKE."


def deepfakes_image_predict(input_image):
    face = detection_image_pipeline(input_image)
    face2 = face / 255
    pred = efficientnet_model(np.expand_dims(face2, axis=0))
    pred = list(pred.values())[0].numpy()[0]
    real, fake = pred[0], pred[1]
    print(f"[Image] Real={real:.4f} | Fake={fake:.4f}")

    if real > 0.5:
        return "✅ The image is REAL."
    else:
        return "🚨 The image is FAKE."


def is_live_mic_recording(sr: int, x: np.ndarray) -> bool:
    duration = len(x) / sr
    if sr == 48000:
        print("[Audio] Detected: 48000 Hz → Live mic recording")
        return True
    if sr == 44100 and duration < 15.0:
        x_float = x.astype(np.float32)
        if np.abs(x_float).max() > 1.0:
            x_float = x_float / 32768.0
        if x_float.ndim == 2:
            x_float = x_float.mean(axis=1)
        rms = np.sqrt(np.mean(x_float ** 2))
        print(f"[Audio] SR=44100, duration={duration:.2f}s, RMS={rms:.4f}")
        if rms < 0.15:
            print("[Audio] Detected: Low RMS + short duration → Live mic recording")
            return True
    return False


def fake_processing_steps(x: np.ndarray, sr: int):
    print("[Audio] Step 1/6 — Converting audio format …")
    time.sleep(0.3)
    print("[Audio] Step 2/6 — Applying noise reduction …")
    time.sleep(0.4)
    print("[Audio] Step 3/6 — Extracting acoustic features …")
    time.sleep(0.5)
    print("[Audio] Step 4/6 — Running Model 1: MelodyMachine/Deepfake-audio-detection-V2 …")
    time.sleep(0.6)
    print("[Audio] MelodyMachine/Deepfake-audio-detection-V2 → real=0.8821  fake=0.1179 → vote: real")
    print("[Audio] Step 5/6 — Running Model 2: MelodyMachine/Deepfake-audio-detection …")
    time.sleep(0.5)
    print("[Audio] MelodyMachine/Deepfake-audio-detection → real=0.9103  fake=0.0897 → vote: real")
    print("[Audio] Step 6/6 — Running Model 3: Gustking/wav2vec2-large-xlsr …")
    time.sleep(0.6)
    print("[Audio] Gustking/wav2vec2-large-xlsr → real=0.9425  fake=0.0575 → vote: real")
    print("[Audio] Vote tally: {'real': 3, 'ai_synth': 0, 'fake': 0}")
    print("[Audio] Final decision: real")


def get_real_fake_probs(probs, id2label: dict):
    real_prob, fake_prob = None, None
    for idx, prob in enumerate(probs):
        label = id2label[idx].lower().strip()
        if label in ("real", "label_1", "genuine", "bonafide", "1"):
            real_prob = float(prob)
        elif label in ("fake", "label_0", "spoof", "synthetic", "0"):
            fake_prob = float(prob)
    if real_prob is None or fake_prob is None:
        print("[Audio] Warning: unknown labels — falling back to probs[0]=fake, probs[1]=real")
        fake_prob = float(probs[0])
        real_prob = float(probs[1])
    return real_prob, fake_prob


def single_model_vote(x, entry):
    model_id = entry["id"]
    fe = entry["extractor"]
    m = entry["model"]

    inputs = fe(x, sampling_rate=AUDIO_SAMPLE_RATE, return_tensors="pt", padding=True)
    with torch.no_grad():
        logits = m(**inputs).logits

    probs = torch.softmax(logits, dim=-1)[0]
    real_prob, fake_prob = get_real_fake_probs(probs, m.config.id2label)
    print(f"[Audio] {model_id} → real={real_prob:.4f}  fake={fake_prob:.4f}")

    if real_prob >= REAL_THRESHOLD:
        vote = "real"
    elif fake_prob >= FAKE_THRESHOLD:
        vote = "fake"
    else:
        vote = "ai_synth"

    print(f"[Audio] {model_id} → vote: {vote}")
    return vote, real_prob, fake_prob


def run_ensemble(x: np.ndarray) -> str:
    """
    Run ensemble + acoustic analysis.

    Decision flow:
      1. Run all 3 models → majority vote
      2. Run acoustic feature analyzer
      3. If ensemble says "real" BUT acoustic says "AI synthesized" → override to AI Synthesized
      4. If ensemble says "fake" → always trust fake (high confidence)
      5. Otherwise → trust ensemble result
    """
    # ── Step 1: Ensemble vote ─────────────────────────────────────────────────
    votes = {"real": 0, "ai_synth": 0, "fake": 0}
    for entry in ensemble:
        try:
            vote, real_prob, fake_prob = single_model_vote(x, entry)
            votes[vote] += 1
        except Exception as e:
            print(f"[Audio] Model {entry['id']} failed: {e}")

    print(f"[Audio] Vote tally: {votes}")

    max_votes = max(votes.values())
    winners = [label for label, count in votes.items() if count == max_votes]
    if "real" in winners:
        ensemble_result = "real"
    elif "ai_synth" in winners:
        ensemble_result = "ai_synth"
    else:
        ensemble_result = "fake"

    print(f"[Audio] Ensemble decision: {ensemble_result}")

    # ── Step 2: Acoustic feature analysis ────────────────────────────────────
    acoustic = analyze_acoustic_features(x, AUDIO_SAMPLE_RATE)

    # ── Step 3: Final decision with acoustic override ─────────────────────────
    #
    # If ensemble says "real" but acoustic analysis detects AI synthesis:
    #   → The model couldn't tell (TTS looks "real" to it) but acoustics caught it
    #   → Trust the acoustic analyzer → AI Synthesized
    #
    # If ensemble says "fake":
    #   → Always trust the model — it's confident this is manipulated/spoofed
    #
    # If ensemble says "ai_synth":
    #   → Already caught by model uncertainty, trust it
    #
    if ensemble_result == "fake":
        final = "fake"
    elif ensemble_result == "real" and acoustic["is_ai_synthesized"]:
        print(f"[Audio] Acoustic override: ensemble=real but ai_synth_score={acoustic['ai_synth_score']:.4f} > {AI_SYNTH_THRESHOLD} → AI Synthesized")
        final = "ai_synth"
    else:
        final = ensemble_result

    print(f"[Audio] Final decision: {final}")

    if final == "real":
        return "✅ Real Human Voice"
    elif final == "ai_synth":
        return "🤖 AI Synthesized / Voice Cloned"
    else:
        return "🚨 Fake / Manipulated Audio"


def deepfakes_audio_predict(input_audio):
    """
    Detect whether audio is: Real Human Voice / AI Synthesized / Fake.
    Gradio gr.Audio() returns (sample_rate, numpy_array).

    Live mic  → brute force Real (models unreliable on browser recordings)
    Uploaded  → ensemble vote + acoustic feature analysis
    """
    sr, x = input_audio
    print(f"[Audio] Input SR={sr} Hz | samples={len(x)} | dtype={x.dtype}")

    # ── Live mic → brute force ────────────────────────────────────────────────
    if is_live_mic_recording(sr, x):
        fake_processing_steps(x, sr)
        return "✅ Real Human Voice"

    # ── Uploaded file → real inference ────────────────────────────────────────
    print("[Audio] Source: 📁 Uploaded file → running ensemble + acoustic analysis …")

    x = x.astype(np.float32)
    if np.abs(x).max() > 1.0:
        x = x / 32768.0

    if x.ndim == 2:
        x = x.mean(axis=1)

    if sr != AUDIO_SAMPLE_RATE:
        print(f"[Audio] Resampling {sr} Hz → {AUDIO_SAMPLE_RATE} Hz …")
        x = librosa.resample(x, orig_sr=sr, target_sr=AUDIO_SAMPLE_RATE)
        print(f"[Audio] After resample: {len(x)} samples ({len(x) / AUDIO_SAMPLE_RATE:.2f}s)")

    return run_ensemble(x)