import gradio as gr
from weight_detector import WeightDetector
import tempfile
import os
from PIL import Image
import requests
from io import BytesIO

# Initialize detector
detector = WeightDetector()

def process_input(image_source: str, image_upload=None, image_url: str = "") -> tuple:
    """
    Process image from different sources (upload, webcam, or URL)
    Returns:
        tuple: (detected_weight, detection_metadata, annotated_image)
    """
    temp_img_path = None
    
    try:
        # Handle different input types
        if image_source == "upload" and image_upload is not None:
            img = image_upload
        elif image_source == "url" and image_url:
            response = requests.get(image_url)
            img = Image.open(BytesIO(response.content))
        else:
            return None, "No valid image provided", None
            
        # Save to temp file for processing
        with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as f:
            temp_img_path = f.name
            img.save(f.name)
            
        # Detect weight
        weight, metadata, annotated_img = detector.detect_weight(temp_img_path)
        
        # Format result message
        if weight is not None:
            message = f"✅ Detected weight: {weight}g"
            if len(metadata) > 1:
                message += f" (from {len(metadata)} possible values)"
        else:
            message = "❌ No weight value detected"
            
        return weight, message, annotated_img
        
    except Exception as e:
        return None, f"Error: {str(e)}", None
    finally:
        if temp_img_path and os.path.exists(temp_img_path):
            os.unlink(temp_img_path)

# Gradio interface
with gr.Blocks(title="Auto Weight Logger") as demo:
    gr.Markdown("""
    # Auto Weight Logger
    Capture or upload an image of a weight measurement to automatically detect and log the value.
    """)
    
    with gr.Row():
        with gr.Column():
            image_source = gr.Radio(
                ["upload", "url"],
                label="Image Source",
                value="upload"
            )
            
            image_upload = gr.Image(
                sources=["upload"],
                type="pil",
                label="Upload Image"
            )
            
            image_url = gr.Textbox(
                label="Image URL",
                visible=False
            )
            
            submit_btn = gr.Button("Detect Weight")
            
        with gr.Column():
            weight_value = gr.Number(
                label="Detected Weight (g)",
                interactive=False
            )
            
            result_message = gr.Textbox(
                label="Detection Result",
                interactive=False
            )
            
            annotated_image = gr.Image(
                label="Annotated Image",
                interactive=False
            )
    
    # Show/hide URL input based on selection
    def toggle_url_visibility(source):
        return gr.Textbox(visible=source == "url")
    
    image_source.change(
        toggle_url_visibility,
        inputs=image_source,
        outputs=image_url
    )
    
    # Process submission
    submit_btn.click(
        process_input,
        inputs=[image_source, image_upload, image_url],
        outputs=[weight_value, result_message, annotated_image]
    )

# For Hugging Face Spaces
demo.launch()