import gradio as gr
import cv2
import numpy as np
from PIL import Image
import logging
from ocr_engine import extract_weight_from_image
from datetime import datetime
import pytz
import sys

# Set up logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s', handlers=[logging.StreamHandler(sys.stdout)])

def process_image(img):
    try:
        # Convert Gradio image (PIL) to process
        if img is None:
            return "No image provided", 0.0, "", None
        
        # Resize if > 5MB
        img_bytes = img.tobytes()
        size_mb = len(img_bytes) / (1024 * 1024)
        if size_mb > 5:
            scale = 0.9
            while size_mb > 5:
                w, h = img.size
                img = img.resize((int(w * scale), int(h * scale)), Image.Resampling.LANCZOS)
                img_bytes = img.tobytes()
                size_mb = len(img_bytes) / (1024 * 1024)
                scale *= 0.9
            logging.info(f"Resized image to {size_mb:.2f} MB")

        # Extract weight
        weight, confidence, unit = extract_weight_from_image(img)
        
        # Return results
        return f"{weight} {unit} (Confidence: {confidence:.2f}%)", f"Processed at {datetime.now(pytz.timezone('Asia/Kolkata')).strftime('%d-%m-%Y %I:%M:%S %p IST')}", img
    except Exception as e:
        logging.error(f"Error in process_image: {str(e)}")
        return f"Error: {str(e)}", "", None

# Gradio interface
with gr.Blocks(title="Auto Weight Logger") as demo:
    gr.Markdown("""
    # 📷 Auto Weight Logger — OCR-Based Smart Scale Reader
    This app detects weight from uploaded or captured images of digital balance displays. Optimized for 7-segment displays and various formats, it extracts numeric weights with high accuracy.
    """)
    
    with gr.Row():
        with gr.Column():
            image_input = gr.Image(source="upload", tool="select", type="pil", label="Upload Weight Display Image")
            webcam_input = gr.Image(source="webcam", type="pil", label="Or Capture with Webcam")
            submit_btn = gr.Button("Detect Weight")
        
        with gr.Column():
            output_text = gr.Textbox(label="Detected Weight", interactive=False)
            timestamp_text = gr.Textbox(label="Processed At", interactive=False)
            output_image = gr.Image(label="Processed Image")

    submit_btn.click(
        fn=process_image,
        inputs=[image_input],
        outputs=[output_text, timestamp_text, output_image]
    )
    
    webcam_input.change(
        fn=process_image,
        inputs=[webcam_input],
        outputs=[output_text, timestamp_text, output_image]
    )

demo.launch()