import os
import warnings

# ==========================================
# 1. ENVIRONMENT CONFIGURATION (Must be at the top)
# ==========================================
# Bypass internal container proxies that block health checks
os.environ["NO_PROXY"] = "localhost,127.0.0.1,::1"
# Force Gradio to bind to all IPs on the correct port
os.environ["GRADIO_SERVER_NAME"] = "0.0.0.0"
os.environ["GRADIO_SERVER_PORT"] = "7860"

# Suppress the non-fatal timm legacy warnings to keep logs clean
warnings.filterwarnings("ignore", category=UserWarning, module="timm.models._factory")

import torch
import gradio as gr
from torchvision import transforms
from model import HybridDeepfakeDetector

# ==========================================
# 2. MODEL INITIALIZATION
# ==========================================
model = HybridDeepfakeDetector()
model.load_state_dict(
    torch.load("deepfake_detector_phase2.pth", 
               map_location="cpu", 
               weights_only=True)
)
model.eval()

transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(
        mean=[0.485, 0.456, 0.406],
        std=[0.229, 0.224, 0.225]
    )
])

# ==========================================
# 3. PREDICTION LOGIC
# ==========================================
def predict(image):
    # Safety check: Prevent crash if user clicks submit before image uploads
    if image is None:
        return "Please upload an image."
        
    tensor = transform(image).unsqueeze(0)
    with torch.no_grad():
        prob = model(tensor).item()
    
    print(f"Raw prob: {prob:.4f}")
    label = "REAL" if prob > 0.5 else "FAKE"
    confidence = prob if label == "REAL" else 1 - prob
    return f"{label} ({confidence*100:.1f}% confident)"

# ==========================================
# 4. UI & DEPLOYMENT
# ==========================================
demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil"),
    outputs=gr.Text(label="Prediction"),
    title="Deepfake Detector",
    description="Upload a face image to detect if it is real or AI-generated."
)

if __name__ == "__main__":
    # Optimal launch parameters for Hugging Face Spaces
    demo.launch(
        share=False,       # Let HF Spaces handle the public URL tunneling
        ssr_mode=False,    # Disable Server-Side Rendering to prevent localhost loop crashes
        show_error=True    # Surface actual code errors to the UI if they happen
    )