import gradio as gr
from datetime import datetime
import pytz
from ocr_engine import extract_weight_from_image

def process_image(img):
    if img is None:
        return "No image uploaded", None, None

    # Get IST time
    ist_time = datetime.now(pytz.timezone("Asia/Kolkata")).strftime("%d-%m-%Y %I:%M:%S %p")
    
    # Run OCR
    weight, confidence = extract_weight_from_image(img)

    # Format output with "kg"
    if weight.replace('.', '', 1).isdigit():
        formatted_weight = f"{weight} kg (Confidence: {confidence}%)"
    else:
        formatted_weight = f"{weight} (Confidence: {confidence}%)"

    return formatted_weight, ist_time, img

# Gradio UI
with gr.Blocks(title="⚖️ Auto Weight Logger") as demo:
    gr.Markdown("# ⚖️ Auto Weight Logger")
    gr.Markdown("Upload or capture an image of a **digital scale display**. It will auto-detect the weight in kg.")

    with gr.Row():
        image_input = gr.Image(type="pil", label="📷 Upload or Capture Image")
        output_weight = gr.Textbox(label="⚖️ Detected Weight (in kg)")

    with gr.Row():
        timestamp = gr.Textbox(label="🕒 Captured At (IST)")
        snapshot = gr.Image(label="📸 Snapshot Image")

    submit = gr.Button("🔍 Detect Weight")
    submit.click(process_image, inputs=image_input, outputs=[output_weight, timestamp, snapshot])

demo.launch()