import os
import gradio as gr
import json
from gradio_client import Client, handle_file

# Validate environment variables and initialize backend client
BACKEND_URL = os.getenv("BACKEND")
HF_TOKEN = os.getenv("TOKEN")

if not BACKEND_URL:
    raise ValueError(
        "BACKEND environment variable is not set. "
        "Please set it to the backend URL (e.g., 'https://your-backend-url')"
    )

try:
    backend = Client(BACKEND_URL, hf_token=HF_TOKEN)
except Exception as e:
    raise Exception(f"Failed to initialize backend client: {str(e)}")

def detect(image):
    """Detect deepfake content in an image with comprehensive error handling"""
    if image is None:
        raise gr.Error("Please upload an image to analyze")
    
    try:
        result_text = backend.predict(
            image=handle_file(image),
            api_name="/detect"
        )
        
        result = json.loads(result_text)
        if not result or result.get("status") != "ok":
            raise gr.Error("Analysis failed: Invalid response from backend")
            
        overall = f"{result['overall']}% Confidence"
        aigen = f"{result['aigen']}% (AI-Generated Content Likelihood)"
        deepfake = f"{result['deepfake']}% (Face Manipulation Likelihood)"
        
        return overall, aigen, deepfake
    
    except json.JSONDecodeError:
        raise gr.Error("Error processing analysis results")
    except Exception as e:
        raise gr.Error(f"Analysis error: {str(e)}")

# [Rest of your CSS and UI code remains the same...]
# I'll include just the essential setup part here for brevity

custom_css = """
.container {
    max-width: 1200px;
    margin: 0 auto;
    padding: 20px;
    font-family: 'Arial', sans-serif;
}
.header {
    color: #2c3e50;
    border-bottom: 2px solid #3498db;
    padding-bottom: 10px;
}
.button-gradient {
    background: linear-gradient(45deg, #3498db, #2ecc71, #9b59b6);
    background-size: 400% 400%;
    border: none;
    padding: 12px 24px;
    font-size: 16px;
    font-weight: 600;
    color: white;
    border-radius: 8px;
    cursor: pointer;
    transition: all 0.3s ease;
    animation: gradientAnimation 3s ease infinite;
    box-shadow: 0 2px 8px rgba(52, 152, 219, 0.3);
}
.button-gradient:hover {
    transform: translateY(-2px);
    box-shadow: 0 4px 12px rgba(52, 152, 219, 0.5);
}
@keyframes gradientAnimation {
    0% { background-position: 0% 50%; }
    50% { background-position: 100% 50%; }
    100% { background-position: 0% 50%; }
}
"""

MARKDOWN0 = """
<div class="header">
    <h1>DeepFake Detection System</h1>
    <p>Advanced AI-powered analysis for identifying manipulated media</p>
</div>
"""

with gr.Blocks(css=custom_css, theme=gr.themes.Default()) as demo:
    gr.Markdown(MARKDOWN0)
    with gr.Row(elem_classes="container"):
        with gr.Column(scale=1):
            image = gr.Image(type='filepath', height=400, label="Upload Image")
            detect_button = gr.Button("Analyze Image", elem_classes="button-gradient")
        with gr.Column(scale=2):
            overall = gr.Label(label="Confidence Score")
            aigen = gr.Label(label="AI-Generated Content")
            deepfake = gr.Label(label="Face Manipulation")
    
    detect_button.click(
        fn=detect,
        inputs=[image],
        outputs=[overall, aigen, deepfake]
    )

demo.queue(api_open=False, concurrency_count=8).launch(
    server_name="0.0.0.0",
    show_api=False,
    debug=True
)