import os
import sys
import gc
import gradio as gr
import torch
import numpy as np
import cv2
import onnxruntime
import spaces
from PIL import Image
from transformers import CLIPImageProcessor
from huggingface_hub import hf_hub_download

# DUMMY GPU FUNCTION (Required for ZeroGPU startup)
@spaces.GPU
def _dummy_gpu_check():
    """
    This function is required to satisfy Hugging Face's ZeroGPU startup check.
    Our actual models run on CPU to save resources, so this is never used.
    """
    return {"status": "ok"}



# CONFIGURATION & SETUP
_DEVICE = "cpu"
FUSION_MODEL_ID = "knmrfr/deepfake-detector"
FUSION_PROCESSOR_ID = "openai/clip-vit-large-patch14"
DETECTRA_REPO = "TensorFold/detectra-v1"

DETECTRA_THRESHOLD = 0.65  
FUSION_THRESHOLD = 0.85    

_face_cascade = None

def _extract_face(image):
    global _face_cascade
    if _face_cascade is None:
        cascade_path = cv2.data.haarcascades + "haarcascade_frontalface_default.xml"
        _face_cascade = cv2.CascadeClassifier(cascade_path)
        
    if image.mode != "RGB": image = image.convert("RGB")
    try:
        img_bgr = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
        gray = cv2.equalizeHist(cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY))
        faces = _face_cascade.detectMultiScale(gray, 1.1, 5, minSize=(30, 30))
        if len(faces) == 0: return None
        x, y, w, h = faces[0]
        return image.crop((x, y, x + w, y + h)).resize((224, 224), Image.BICUBIC)
    except Exception:
        return None

def _run_detectra(image):
    session = None
    try:
        detectra_path = hf_hub_download(repo_id=DETECTRA_REPO, filename="model.onnx")
        session = onnxruntime.InferenceSession(detectra_path, providers=['CPUExecutionProvider'])
        input_name = session.get_inputs()[0].name

        w, h = image.size
        scale = 440 / min(w, h)
        new_w, new_h = int(w * scale), int(h * scale)
        img_resized = image.resize((new_w, new_h), Image.BILINEAR)
        left, top = (new_w - 384) // 2, (new_h - 384) // 2
        img_cropped = img_resized.crop((left, top, left + 384, top + 384))
        
        img_array = np.array(img_cropped).astype(np.float32) / 255.0
        mean = np.array([0.485, 0.456, 0.406], dtype=np.float32)
        std = np.array([0.229, 0.224, 0.225], dtype=np.float32)
        img_array = (img_array - mean) / std
        img_array = np.transpose(img_array, (2, 0, 1))
        img_array = np.expand_dims(img_array, axis=0)
        
        logit = session.run(None, {input_name: img_array})[0]
        return float(1 / (1 + np.exp(-logit[0][0])))
    except Exception as e:
        print(f"Detectra error: {e}")
        return 0.0
    finally:
        if session is not None: del session
        gc.collect()

def _run_fusion(image):
    model = None
    processor = None
    try:
        model_code_path = hf_hub_download(repo_id=FUSION_MODEL_ID, filename="model.py")
        sys.path.insert(0, os.path.dirname(model_code_path))
        from model import DeepfakeDetector
        
        model = DeepfakeDetector.from_pretrained(FUSION_MODEL_ID).eval().to(_DEVICE)
        processor = CLIPImageProcessor.from_pretrained(FUSION_PROCESSOR_ID)
        
        face = _extract_face(image)
        if face is None: return 0.0

        inputs = processor(images=face, return_tensors="pt").to(_DEVICE)
        with torch.no_grad():
            outputs = model(**inputs)
            logits = outputs.logits if hasattr(outputs, "logits") else outputs["logits"]
            probs = torch.softmax(logits, dim=1)[0]
            
        fake_idx = 1
        for i, label in model.id2label.items():
            if "fake" in str(label).lower(): fake_idx = int(i); break
        return probs[fake_idx].item()
    except Exception as e:
        print(f"Fusion error: {e}")
        return 0.0
    finally:
        if model is not None: del model
        if processor is not None: del processor
        gc.collect()

def detect_image(image):
    if image is None:
        return "No image provided", 0.0, {}
        
    image = image.convert("RGB")
    ai_prob = _run_detectra(image)
    fusion_prob = _run_fusion(image)
    
    if ai_prob > DETECTRA_THRESHOLD:
        verdict = "AI-Generated"
        confidence = ai_prob * 100
    elif fusion_prob > FUSION_THRESHOLD:
        verdict = "Deepfake (Manipulated)"
        confidence = fusion_prob * 100
    else:
        verdict = "Real"
        confidence = ((1 - ai_prob) + (1 - fusion_prob)) / 2 * 100

    scores = {
        "AI_Score": round(ai_prob * 100, 2),
        "Deepfake_Score": round(fusion_prob * 100, 2)
    }
    
    return verdict, round(confidence, 2), scores


# GRADIO INTERFACE & LAUNCH
demo = gr.Interface(
    fn=detect_image,
    inputs=gr.Image(type="pil", label="Upload Image"),
    outputs=[
        gr.Textbox(label="Verdict"),
        gr.Number(label="Confidence (%)"),
        gr.JSON(label="Detailed Scores")
    ],
    title="Deepfake & AI Image Detector",
    description="Upload an image to check if it's Real, AI-Generated, or Manipulated."
)

# Launch at module level so Hugging Face detects the app correctly
if __name__ == "__main__":
    demo.launch()