import gradio as gr
from gradio_client import Client
import os
import tempfile

BACKEND_URL = os.environ.get("BACKEND_URL", "").strip()

try:
    client = Client(BACKEND_URL, headers={"ngrok-skip-browser-warning": "true"})
    backend_available = True
except:
    client = None
    backend_available = False

def process_media(file_obj, webcam_img, model_type, conf_thresh, max_dets, task_type):
    """Process media - backend expects both file and webcam paths"""
    if not client:
        return [gr.update()] * 5
    
    try:
        # Convert webcam PIL to file path if present
        webcam_path = None
        if webcam_img is not None:
            with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as tmp:
                webcam_img.save(tmp, 'PNG')
                webcam_path = tmp.name
        
        # Backend expects both parameters - use None for missing one
        result = client.predict(
            uploaded_file_obj=file_obj if file_obj else None,
            webcam_image_pil=webcam_path if webcam_path else None,
            model_type_choice=model_type,
            conf_threshold_ui=conf_thresh,
            max_detections_ui=max_dets,
            task_type=task_type,
            api_name="/process_media"
        )
        
        # Cleanup
        if webcam_path and os.path.exists(webcam_path):
            os.unlink(webcam_path)
        
        return result
        
    except Exception as e:
        print(f"Process error: {e}")
        # Return error message in processed image slot
        return [
            gr.update(),  # raw image file
            gr.update(),  # raw video file  
            gr.update(),  # raw image webcam
            gr.update(value=None, visible=True),  # processed image - show error
            gr.update()   # processed video
        ]

# Simplified interface without complex preview forwarding
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🐵 PrimateFace Detection, Pose & Gaze Demo")
    
    if not backend_available:
        gr.Markdown("### 🔴 GPU Server Offline")
    else:
        with gr.Row():
            with gr.Column():
                with gr.Tabs():
                    with gr.TabItem("Upload"):
                        input_file = gr.File(label="Upload Image/Video")
                        # Simple local preview
                        preview_img = gr.Image(label="Preview", visible=False)
                    
                    with gr.TabItem("Webcam"):
                        input_webcam = gr.Image(sources=["webcam"], type="pil")
                
                clear_btn = gr.Button("Clear All")
            
            with gr.Column():
                gr.Markdown("### Results")
                output_image = gr.Image(label="Processed", visible=False)
                output_video = gr.Video(label="Processed", visible=False)
        
        # Examples
        gr.Examples(
            examples=[["images/" + f] for f in [
                "allocebus_000003.jpeg",
                "tarsius_000120.jpeg", 
                "nasalis_proboscis-monkey.png",
                "macaca_000032.jpeg",
                "mandrillus_000011.jpeg",
                "pongo_000006.jpeg"
            ]],
            inputs=input_file
        )
        
        submit_btn = gr.Button("Detect Faces", variant="primary")
        
        # Controls
        model_choice = gr.Radio(["MMDetection"], value="MMDetection", visible=False)
        task_type = gr.Dropdown(
            ["Face Detection", "Face Pose Estimation", "Gaze Estimation [experimental]"],
            value="Face Detection"
        )
        conf_threshold = gr.Slider(0.05, 0.95, 0.25, step=0.05, label="Confidence")
        max_detections = gr.Slider(1, 10, 3, step=1, label="Max Detections")
        
        # Simple local preview for uploaded files
        def show_preview(file):
            if file and file.name.lower().endswith(('.jpg', '.jpeg', '.png', '.bmp')):
                return gr.update(value=file, visible=True)
            return gr.update(visible=False)
        
        input_file.change(show_preview, inputs=[input_file], outputs=[preview_img])
        
        # Main processing - only use last 3 outputs (skip raw previews)
        def process_and_extract_outputs(*args):
            result = process_media(*args)
            # Return only processed outputs
            return result[-2:]  # Just processed image and video
        
        submit_btn.click(
            process_and_extract_outputs,
            inputs=[input_file, input_webcam, model_choice, conf_threshold, max_detections, task_type],
            outputs=[output_image, output_video]
        )
        
        # Simple clear
        clear_btn.click(
            lambda: [gr.update(value=None), gr.update(value=None), gr.update(visible=False), 
                    gr.update(visible=False), gr.update(visible=False)],
            outputs=[input_file, input_webcam, preview_img, output_image, output_video]
        )

demo.launch()