import os, io, tempfile, warnings
import numpy as np
import gradio as gr

# =========================
# TensorFlow / Keras Setup
# =========================
import tensorflow as tf
from tensorflow.keras.applications.resnet50 import preprocess_input as resnet_preprocess

# Limit TensorFlow thread usage (prevents OOM)
os.environ["OMP_NUM_THREADS"] = "2"
os.environ["MKL_NUM_THREADS"] = "2"
os.environ["OPENBLAS_NUM_THREADS"] = "2"
os.environ["NUMEXPR_NUM_THREADS"] = "2"
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"

# Try to enable GPU safely for TensorFlow
try:
    gpus = tf.config.experimental.list_physical_devices("GPU")
    if gpus:
        for gpu in gpus:
            tf.config.experimental.set_memory_growth(gpu, True)
        tf.keras.mixed_precision.set_global_policy("mixed_float16")
        print("[INFO] ✅ TensorFlow GPU detected & configured (mixed precision ON).")
    else:
        print("[WARN] ⚠️ No TensorFlow GPU detected, running on CPU.")
except Exception as e:
    print(f"[WARN] TensorFlow GPU init skipped: {e}")

# =========================
# PyTorch / MTCNN Setup
# =========================
import torch
from facenet_pytorch import MTCNN
import cv2
import dlib
from imutils import face_utils
from scipy.spatial import distance as dist

warnings.filterwarnings("ignore")

# Use GPU for MTCNN if available
if torch.cuda.is_available():
    _torch_device = torch.device("cuda")
    torch.backends.cudnn.benchmark = True
    try:
        torch.set_float32_matmul_precision("medium")
    except Exception:
        pass
    print("[INFO] ✅ PyTorch GPU detected — MTCNN will use CUDA.")
else:
    _torch_device = torch.device("cpu")
    print("[WARN] ⚠️ No GPU detected — MTCNN running on CPU.")

# =========================
# Paths / Config
# =========================
VIDEO_MODEL_PATH = "models/video_model.h5"
DLIB_LANDMARK_MODEL = "shape_predictor_68_face_landmarks.dat"

IMG_SIZE = (224, 224)
FRAME_STEP = 5
NUM_MAX_FACES = 300
EAR_THRESHOLD = 0.25
EAR_CONSEC_FRAMES = 3
PRED_THRESHOLD = 0.5

# Lazy-loaded globals
_video_model = None
_mtcnn = None
_dlib_detector = None
_dlib_predictor = None


def lazy_load():
    global _video_model, _mtcnn, _dlib_detector, _dlib_predictor

    if _video_model is None:
        if not os.path.exists(VIDEO_MODEL_PATH):
            raise FileNotFoundError(f"Missing: {VIDEO_MODEL_PATH}")
        # Allow tf to place on GPU if available
        _video_model = tf.keras.models.load_model(VIDEO_MODEL_PATH)
        print("[INFO] Video model loaded.")

    if _mtcnn is None:
        _mtcnn = MTCNN(keep_all=True, device=_torch_device, image_size=IMG_SIZE[0])
        print(f"[INFO] MTCNN ready on {_torch_device}.")

    if _dlib_detector is None or _dlib_predictor is None:
        if not os.path.exists(DLIB_LANDMARK_MODEL):
            raise FileNotFoundError(
                f"Missing dlib predictor: {DLIB_LANDMARK_MODEL}. Place it beside app.py."
            )
        _dlib_detector = dlib.get_frontal_face_detector()
        _dlib_predictor = dlib.shape_predictor(DLIB_LANDMARK_MODEL)
        print("[INFO] dlib detector + predictor ready.")

# =========================
# VIDEO: faces + blink features
# =========================
def _eye_aspect_ratio(eye_pts):
    A = dist.euclidean(eye_pts[1], eye_pts[5])
    B = dist.euclidean(eye_pts[2], eye_pts[4])
    C = dist.euclidean(eye_pts[0], eye_pts[3])
    return (A + B) / (2.0 * C + 1e-9)

def extract_faces_all(path):
    """Sample every 5th frame, keep ALL faces per frame, resize to 224x224."""
    cap = cv2.VideoCapture(path)
    total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    if total <= 0:
        cap.release()
        return None

    faces = []
    for idx in range(0, total, FRAME_STEP):
        if len(faces) >= NUM_MAX_FACES:
            break
        cap.set(cv2.CAP_PROP_POS_FRAMES, idx)
        ok, frame = cap.read()
        if not ok:
            break
        rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        boxes, _ = _mtcnn.detect(rgb)
        if boxes is None:
            continue
        for (x1, y1, x2, y2) in boxes.astype(int):
            x1, y1 = max(0, x1), max(0, y1)
            x2, y2 = min(frame.shape[1], x2), min(frame.shape[0], y2)
            crop = frame[y1:y2, x1:x2]
            if crop.size > 0:
                faces.append(cv2.resize(crop, IMG_SIZE))
            if len(faces) >= NUM_MAX_FACES:
                break
    cap.release()
    if not faces:
        return None
    return np.stack(faces, axis=0).astype(np.uint8)

def blink_features_from_crops(faces):
    """Compute EAR on cropped faces; denominator = #frames with valid EAR."""
    (lStart, lEnd) = face_utils.FACIAL_LANDMARKS_IDXS["left_eye"]
    (rStart, rEnd) = face_utils.FACIAL_LANDMARKS_IDXS["right_eye"]

    ear_values = []
    blink_counter = 0
    total_blinks = 0

    for face in faces:
        gray = cv2.cvtColor(face, cv2.COLOR_BGR2GRAY)
        rects = _dlib_detector(gray, 0)
        if len(rects) == 0:
            continue
        shape = _dlib_predictor(gray, rects[0])
        shape = face_utils.shape_to_np(shape)

        left_eye = shape[lStart:lEnd]
        right_eye = shape[rStart:rEnd]
        ear = 0.5 * (_eye_aspect_ratio(left_eye) + _eye_aspect_ratio(right_eye))
        ear_values.append(ear)

        if ear < EAR_THRESHOLD:
            blink_counter += 1
        else:
            if blink_counter >= EAR_CONSEC_FRAMES:
                total_blinks += 1
            blink_counter = 0

    n_ear = len(ear_values)
    if n_ear == 0:
        return np.array([0, 0.0, 0.0], dtype=np.float32)

    blink_freq = total_blinks / float(n_ear)
    ear_var = float(np.var(np.array(ear_values, dtype=np.float32)))
    return np.array([total_blinks, blink_freq, ear_var], dtype=np.float32)

def predict_video_prob(video_path):
    faces = extract_faces_all(video_path)
    if faces is None or len(faces) == 0:
        return None, "No faces detected."

    blink_feats = blink_features_from_crops(faces)
    # Tile to per-face samples
    tiled = np.tile(blink_feats, (faces.shape[0], 1)).astype(np.float32)
    imgs = resnet_preprocess(faces.astype(np.float32))

    # Batch size heuristic (bigger if GPU is present)
    has_tf_gpu = len(tf.config.experimental.list_physical_devices("GPU")) > 0
    bs = 128 if has_tf_gpu else 32

    preds = []
    for i in range(0, imgs.shape[0], bs):
        p = _video_model.predict([imgs[i:i+bs], tiled[i:i+bs]], verbose=0)
        preds.append(p.reshape(-1))
    return float(np.mean(np.concatenate(preds, axis=0))), None

def to_verdict(score):
    return "DEEPFAKE" if score >= PRED_THRESHOLD else "REAL"

# =========================
# Inference entry
# =========================
def run_inference(video_file):
    lazy_load()
    if video_file is None:
        return None, None, "Please upload a video file."

    video_prob, vmsg = predict_video_prob(video_file)
    if video_prob is None:
        return None, None, vmsg or "Unable to process the video."

    verdict = f"VIDEO ONLY: {to_verdict(video_prob)}"
    fmt = lambda x: None if x is None else round(float(x), 4)
    return fmt(video_prob), verdict, None

# =========================
# Gradio UI
# =========================
with gr.Blocks(title="Deepfake Detector — Video Only (GPU-ready)") as demo:
    gr.Markdown(
        "### 🎭 Deepfake Detector — **Video Only**\n"
        "- **Compute**: Uses GPU automatically if available (PyTorch MTCNN, TensorFlow model; mixed precision on TF).\n"
        "- **Video path**: samples every 5th frame, keeps **all** faces, resizes to 224×224.\n"
        "- **Features**: ResNet50 preprocessing + blink EAR features on cropped faces via dlib.\n"
        f"- **Threshold**: {PRED_THRESHOLD} (≥ means DEEPFAKE).\n"
    )
    with gr.Row():
        video_in = gr.Video(label="Video")
    go = gr.Button("Analyze")
    with gr.Row():
        v_out = gr.Number(label="Video probability (deepfake)", precision=4)
        verdict_out = gr.Textbox(label="Verdict", interactive=False)
    msg_out = gr.Textbox(label="Message / Warnings", interactive=False)

    go.click(run_inference, inputs=[video_in], outputs=[v_out, verdict_out, msg_out])

if __name__ == "__main__":
    lazy_load()
    demo.launch(server_name="0.0.0.0", server_port=7860)