# Deepfake Detection Gradio App - v1.1
import gradio as gr
import os
import sys
import json
import argparse
from types import SimpleNamespace

# Try to import detector - if this fails, we'll show an error in the UI
try:
    from support.detect import run_detect
    DETECTOR_AVAILABLE = True
    IMPORT_ERROR = None
except Exception as e:
    DETECTOR_AVAILABLE = False
    IMPORT_ERROR = str(e)
    print(f"Warning: Could not import detector: {e}")
    # Create a dummy function
    def run_detect(args):
        raise ImportError(f"Detector not available: {IMPORT_ERROR}")

# Download weights on first run (for HF Spaces)
if os.environ.get("SPACE_ID"):
    try:
        from download_weights import download_all_weights
        download_all_weights()
    except Exception as e:
        print(f"Warning: Could not download weights: {e}")

# Available detectors based on launcher.py
DETECTORS = ['R50_TF', 'R50_nodown', 'CLIP-D', 'P2G', 'NPR']

def predict(image_path, detector_name):
    # Check if detector is available
    if not DETECTOR_AVAILABLE:
        return json.dumps({
            "error": "Detector module not available",
            "details": IMPORT_ERROR,
            "message": "The detection system could not be initialized. Please check the logs."
        }, indent=2)
    
    if not image_path:
        return json.dumps({"error": "Please upload an image."}, indent=2)
    
    # Create a temporary output file path
    output_path = "temp_result.json"
    
    # Mock args object
    args = SimpleNamespace(
        image=image_path,
        detector=detector_name,
        config_dir='configs',
        output=output_path,
        weights='pretrained', # Use default/pretrained
        device='cpu', # Force CPU
        dry_run=False,
        verbose=False
    )
    
    try:
        # Run detection
        # We need to capture stdout/stderr or just trust the function
        # run_detect might raise FileNotFoundError if weights are missing
        run_detect(args)
        
        # Read results
        if os.path.exists(output_path):
            with open(output_path, 'r') as f:
                result = json.load(f)
            
            # Format output
            prediction = result.get('prediction', 'Unknown')
            confidence = result.get('confidence', 0.0)
            elapsed_time = result.get('elapsed_time', 0.0)
            
            output = {
                "Prediction": prediction,
                "Confidence": f"{confidence:.4f}",
                "Elapsed Time": f"{elapsed_time:.3f}s"
            }
            return json.dumps(output, indent=2)
        else:
            return json.dumps({"error": "No result file generated. Check console logs for details."}, indent=2)
            
    except FileNotFoundError as e:
        return json.dumps({
            "error": str(e), 
            "message": f"Please ensure you have downloaded the weights for {detector_name}."
        }, indent=2)
    except Exception as e:
        return json.dumps({"error": str(e)}, indent=2)
    finally:
        # Cleanup
        if os.path.exists(output_path):
            os.remove(output_path)

# Create Gradio Interface
# Use theme only if gradio version supports it
demo = gr.Blocks(title="Deepfake Detection", theme=gr.themes.Soft())

with demo:
    gr.Markdown("# 🔍 Deepfake Detection Library")
    gr.Markdown("""
    Upload an image and select a detector to check if it's real or fake.
    
    **Available Detectors:**
    - **R50_TF**: ResNet-50 based detector
    - **R50_nodown**: ResNet-50 without downsampling
    - **CLIP-D**: CLIP-based detector
    - **P2G**: Prompt2Guard detector
    - **NPR**: Neural Posterior Regularization
    """)
    
    with gr.Row():
        with gr.Column():
            image_input = gr.Image(type="filepath", label="Input Image", height=400)
            detector_input = gr.Dropdown(
                choices=DETECTORS, 
                value=DETECTORS[0], 
                label="Select Detector",
                info="Choose which deepfake detection model to use"
            )
            submit_btn = gr.Button("🔍 Detect", variant="primary")
        
        with gr.Column():
            output_display = gr.Textbox(
                label="Detection Results", 
                lines=15, 
                max_lines=20,
                show_copy_button=True
            )
    
    gr.Markdown("""
    ---
    ### About
    This Space provides access to multiple state-of-the-art deepfake detection models. 
    All models are trained on StyleGAN2, StableDiffusionXL, FFHQ, and FORLAB datasets.
    
    **Note:** First detection may be slower due to model loading.
    """)
    
    submit_btn.click(
        fn=predict,
        inputs=[image_input, detector_input],
        outputs=output_display
    )

if __name__ == "__main__":
    # For HF Spaces, configure server settings
    if os.environ.get("SPACE_ID"):
        demo.launch(server_name="0.0.0.0", server_port=7860, allowed_paths=["."])
    else:
        # Local execution
        demo.launch()