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.")


# ─────────────────────────────────────────────────────────────────────────────
# TEXT DETECTOR: HybridAITextDetector (DeBERTa + BiLSTM + CNN + Transformer)
# ─────────────────────────────────────────────────────────────────────────────
TEXT_CHECKPOINT = "best_text_detector.pt"
TEXT_THRESHOLD  = 0.5   # update with optimal F1 threshold from your training run

_text_detector = None   # lazy-loaded on first call

def _get_text_detector():
    """Lazy-load the text detector (avoids startup delay if tab isn't used)."""
    global _text_detector
    if _text_detector is None:
        from text_detector_inference import TextDetectorInference
        print("[Text] Loading HybridAITextDetector ...")
        _text_detector = TextDetectorInference(
            checkpoint=TEXT_CHECKPOINT,
            threshold=TEXT_THRESHOLD,
        )
        print("[Text] ✅ Text detector ready")
    return _text_detector


# ─────────────────────────────────────────────────────────────────────────────
# ACOUSTIC FEATURE ANALYZER
# ─────────────────────────────────────────────────────────────────────────────
AI_SYNTH_THRESHOLD = 0.60


def analyze_acoustic_features(x: np.ndarray, sr: int) -> dict:
    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
    rms_cv = np.sqrt(rms_variance) / rms_mean
    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}")

    spec_flatness = librosa.feature.spectral_flatness(y=x, hop_length=hop_length)[0]
    mean_flatness = np.mean(spec_flatness)
    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}")

    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)
            pitch_synth_score = max(0.0, min(1.0, 1.0 - (pitch_variance / 0.15)))
        else:
            pitch_synth_score = 0.5
    except Exception:
        pitch_synth_score = 0.5
    print(f"[Acoustic] Pitch variance score={pitch_synth_score:.4f}")

    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
    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}")

    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:
    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}")

    acoustic = analyze_acoustic_features(x, AUDIO_SAMPLE_RATE)

    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):
    sr, x = input_audio
    print(f"[Audio] Input SR={sr} Hz | samples={len(x)} | dtype={x.dtype}")

    if is_live_mic_recording(sr, x):
        fake_processing_steps(x, sr)
        return "✅ Real Human Voice"

    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)


# ─────────────────────────────────────────────────────────────────────────────
# TEXT DEEPFAKE DETECTION
# Hybrid DeBERTa-v3-small + BiLSTM + CNN + Transformer
# Returns: "✅ Human-Written" / "🤖 AI-Generated"
# ─────────────────────────────────────────────────────────────────────────────

def deepfakes_text_predict(input_text: str) -> str:
    """
    Detect whether the input text is human-written or AI-generated.

    Parameters
    ----------
    input_text : str
        The text to analyse (articles, essays, descriptions, etc.)

    Returns
    -------
    str
        A formatted result string for display in the Gradio textbox.
    """
    if not input_text or not input_text.strip():
        return "⚠️ Please enter some text to analyse."

    text = input_text.strip()
    word_count = len(text.split())
    print(f"[Text] Input: {word_count} words")

    if word_count < 10:
        return (
            "⚠️ Input too short — please provide at least 10 words for a reliable result.\n"
            f"   (You entered {word_count} word{'s' if word_count != 1 else ''})"
        )

    try:
        detector = _get_text_detector()
        result   = detector.predict(text)

        if "error" in result:
            return f"❌ Error: {result['error']}"

        label      = result["label"]
        ai_prob    = result["ai_prob"]
        human_prob = result["human_prob"]
        confidence = result["confidence"]

        print(f"[Text] label={label} | ai_prob={ai_prob:.4f} | human_prob={human_prob:.4f}")

        # ── Format output ─────────────────────────────────────────────────────
        if label == "AI-Generated":
            verdict_icon = "🤖"
            verdict_text = "AI-Generated Text"
        else:
            verdict_icon = "✅"
            verdict_text = "Human-Written Text"

        # Confidence bar (ASCII, 20 chars)
        bar_filled = round(confidence * 20)
        bar = "█" * bar_filled + "░" * (20 - bar_filled)

        output = (
            f"{verdict_icon}  {verdict_text}\n"
            f"\n"
            f"Confidence  [{bar}]  {confidence*100:.1f}%\n"
            f"\n"
            f"P(AI-Generated)  : {ai_prob*100:.1f}%\n"
            f"P(Human-Written) : {human_prob*100:.1f}%\n"
            f"\n"
            f"Words analysed   : {word_count}\n"
            f"(First 128 tokens used — ~100 words)"
        )
        return output

    except Exception as e:
        print(f"[Text] ❌ Prediction failed: {e}")
        return f"❌ Text detection failed: {str(e)}\nMake sure best_text_detector.pt is present in the Space."