import gradio as gr
from weight_detector import WeightDetector
import tempfile
import os

detector = WeightDetector()

def process_input(image_source: str, image_upload=None, image_url: str = "") -> dict:
    """Process webcam/image and return weight + IST time"""
    temp_img_path = None
    try:
        # Handle webcam/image upload
        if image_source == "webcam" and image_upload is not None:
            img = image_upload
        elif image_source == "upload" and image_upload is not None:
            img = image_upload
        elif image_source == "url" and image_url:
            import requests
            from io import BytesIO
            response = requests.get(image_url)
            img = Image.open(BytesIO(response.content))
        else:
            return {
                "weight": None,
                "message": "⚠️ No image provided!",
                "image": None,
                "time": detector.get_current_ist()
            }
        
        # Save to temp file
        with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as f:
            temp_img_path = f.name
            img.save(f.name)
        
        # Detect weight
        return detector.detect_weight(temp_img_path)
    
    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.remove(temp_img_path)

# Gradio UI
with gr.Blocks(title="Auto Weight Logger") as demo:
    gr.Markdown("""
    # **⚖️ Auto Weight Logger (7-Segment OCR)**
    **Capture weight from digital balances using a webcam or image upload.**
    - ✅ Optimized for **7-segment displays** (e.g., lab balances)
    - 📅 Logs **IST time** automatically
    - 🚫 Detects **blurry/glare** images
    """)
    
    with gr.Row():
        with gr.Column():
            image_source = gr.Radio(
                ["webcam", "upload", "url"],
                label="Input Source",
                value="webcam"
            )
            
            image_upload = gr.Image(
                sources=["webcam", "upload"],
                type="pil",
                label="Capture/Upload Image",
                interactive=True
            )
            
            image_url = gr.Textbox(
                label="Image URL (if selected)",
                visible=False
            )
            
            submit_btn = gr.Button("Detect Weight", variant="primary")
        
        with gr.Column():
            weight_value = gr.Number(
                label="Detected Weight (g)",
                interactive=False
            )
            
            detection_time = gr.Textbox(
                label="Detection Time (IST)",
                interactive=False
            )
            
            result_message = gr.Textbox(
                label="Result",
                interactive=False
            )
            
            annotated_image = gr.Image(
                label="Annotated Image",
                interactive=False
            )
    
    # Show/hide URL input
    def toggle_url_visibility(source):
        return gr.Textbox(visible=source == "url")
    
    image_source.change(
        toggle_url_visibility,
        inputs=image_source,
        outputs=image_url
    )
    
    # Process input
    submit_btn.click(
        process_input,
        inputs=[image_source, image_upload, image_url],
        outputs=[weight_value, detection_time, result_message, annotated_image]
    )

demo.launch()