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 = "") -> dict:
    """Process image and return results with IST"""
    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 {
                "weight": None,
                "message": "No valid image provided",
                "image": None,
                "time": detector.get_current_ist()
            }
            
        # 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, time, annotated_img = detector.detect_weight(temp_img_path)
        
        # Format result message
        if weight is not None:
            message = f"✅ Detected weight: {weight:.2f}g at {time}"
        else:
            message = f"❌ No weight value detected at {time}"
            
        return {
            "weight": weight,
            "message": message,
            "image": annotated_img,
            "time": time
        }
        
    except Exception as e:
        return {
            "weight": None,
            "message": f"Error: {str(e)}",
            "image": None,
            "time": detector.get_current_ist()
        }
    finally:
        if temp_img_path and os.path.exists(temp_img_path):
            os.unlink(temp_img_path)

# Custom CSS for better mobile display
css = """
#mobile-view {
    display: none;
}
@media screen and (max-width: 768px) {
    #desktop-view {
        display: none;
    }
    #mobile-view {
        display: block;
    }
}
"""

# Gradio interface
with gr.Blocks(title="Auto Weight Logger", css=css) as demo:
    gr.Markdown("""
    # 🏋️ Auto Weight Logger
    Capture or upload an image of a digital scale to automatically detect the weight value.
    """)
    
    with gr.Row():
        with gr.Column():
            image_source = gr.Radio(
                ["upload", "url"],
                label="Image Source",
                value="upload",
                elem_id="source-select"
            )
            
            image_upload = gr.Image(
                sources=["upload", "webcam"],
                type="pil",
                label="Upload Image or Use Webcam",
                elem_id="image-upload"
            )
            
            image_url = gr.Textbox(
                label="Image URL",
                visible=False,
                elem_id="image-url"
            )
            
            submit_btn = gr.Button("Detect Weight", variant="primary")
            
        with gr.Column():
            weight_value = gr.Number(
                label="Detected Weight (grams)",
                interactive=False,
                elem_id="weight-value"
            )
            
            detection_time = gr.Textbox(
                label="Detection Time (IST)",
                interactive=False,
                elem_id="detection-time"
            )
            
            result_message = gr.Textbox(
                label="Result",
                interactive=False,
                elem_id="result-message"
            )
            
            annotated_image = gr.Image(
                label="Annotated Image",
                interactive=False,
                elem_id="annotated-image"
            )
    
    # Mobile view toggle
    with gr.Column(visible=False, elem_id="mobile-view"):
        gr.Markdown("### Mobile Instructions")
        gr.Markdown("1. Tap 'Webcam' to capture\n2. Tap 'Detect Weight'")
    
    # 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": weight_value,
            "message": result_message,
            "image": annotated_image,
            "time": detection_time
        }
    )

# For Hugging Face Spaces
demo.launch()