# 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 Fake": 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 Space", theme=gr.themes.Soft())

with demo:
    gr.Markdown("# 🔍 Deepfake Detection Space")
    gr.Markdown("""
    This space collects a series of state-of-the-art methods for deepfake detection, allowing for free and unlimited use.
    
    ### Training & Performance
    All methods have been trained using the **[DeepShield dataset](https://zenodo.org/records/15648378)**, which includes images generated with **Stable Diffusion XL** and **StyleGAN 2**. 
    You can expect performance comparable to the results shown in [Dell'Anna et al. (2025)](https://arxiv.org/pdf/2504.20658).

    ### Understanding the Results
    *   **Prediction**: Tells if an image is **Real** or **Fake**.
    *   **Confidence Fake**: The confidence with which the model determines if the image is fake.
    *   **Elapsed Time**: The time the model needed to make the prediction (excluding preprocessing or model building).
    """)
    
    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
            )
    
    with gr.Accordion("📚 Model Details", open=False):
        gr.Markdown("""
        ### **R50_TF**
        *   **Description**: A ResNet50 architecture modified to exclude downsampling at the first layer. It uses "learned prototypes" in the classification head for robust detection.
        *   **Paper**: [TrueFake: A Real World Case Dataset of Last Generation Fake Images also Shared on Social Networks](https://arxiv.org/pdf/2504.20658)
        *   **Code**: [GitHub Repository](https://github.com/MMLab-unitn/TrueFake-IJCNN25)

        ### **R50_nodown**
        *   **Description**: A ResNet-50 model without downsampling operations in the first layer, designed to preserve high-frequency artifacts common in synthetic images.
        *   **Paper**: [On the detection of synthetic images generated by diffusion models](https://arxiv.org/abs/2211.00680)
        *   **Code**: [GitHub Repository](https://grip-unina.github.io/DMimageDetection/)

        ### **CLIP-D**
        *   **Description**: A lightweight detection strategy based on CLIP features. It exhibits surprising generalization ability using only a handful of example images.
        *   **Paper**: [Raising the Bar of AI-generated Image Detection with CLIP](https://arxiv.org/abs/2312.00195v2)
        *   **Code**: [GitHub Repository](https://grip-unina.github.io/ClipBased-SyntheticImageDetection/)

        ### **P2G (Prompt2Guard)**
        *   **Description**: Uses Vision-Language Models (VLMs) with conditioned prompt-optimization for continual deepfake detection. It leverages read-only prompts for efficiency.
        *   **Paper**: [Conditioned Prompt-Optimization for Continual Deepfake Detection](https://arxiv.org/abs/2407.21554)
        *   **Code**: [GitHub Repository](https://github.com/laitifranz/Prompt2Guard)

        ### **NPR**
        *   **Description**: Focuses on Neighboring Pixel Relationships (NPR) to capture generalized structural artifacts stemming from up-sampling operations in generative networks.
        *   **Paper**: [Rethinking the Up-Sampling Operations in CNN-based Generative Network for Generalizable Deepfake Detection](https://arxiv.org/abs/2312.10461)
        *   **Code**: [GitHub Repository](https://github.com/chuangchuangtan/NPR-DeepfakeDetection)
        """)

    gr.Markdown("""
    ---
    ### References
    1. Dell'Anna, S., Montibeller, A., & Boato, G. (2025). *TrueFake: A Real World Case Dataset of Last Generation Fake Images also Shared on Social Networks*. arXiv preprint arXiv:2504.20658.
    2. Corvi, R., et al. (2023). *On the detection of synthetic images generated by diffusion models*. ICASSP.
    3. Cozzolino, D., et al. (2023). *Raising the Bar of AI-generated Image Detection with CLIP*. CVPRW.
    4. Laiti, F., et al. (2024). *Conditioned Prompt-Optimization for Continual Deepfake Detection*. arXiv preprint arXiv:2407.21554.
    5. Tan, C., et al. (2024). *Rethinking the up-sampling operations in cnn-based generative network for generalizable deepfake detection*. CVPR.
    """)
    
    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()