import gradio as gr
from PIL import Image, UnidentifiedImageError
from datetime import datetime
import pytz
from ocr_engine import extract_weight_from_image  # Make sure this file exists

def detect_weight(image: Image.Image):
    try:
        # Get IST timestamp
        ist_time = datetime.now(pytz.timezone("Asia/Kolkata")).strftime("%Y-%m-%d %H:%M:%S")

        # Perform OCR
        weight, confidence = extract_weight_from_image(image)

        if confidence > 0:
            result = f"✅ **Detected Weight:** {weight}  \n📊 **Confidence:** {confidence:.2f}  \n🕒 **Captured At (IST):** {ist_time}"
        else:
            result = "❌ No weight detected. Please upload a clearer image."
        return image, result
    except UnidentifiedImageError:
        return None, "❌ Invalid image format."

# Gradio UI
demo = gr.Interface(
    fn=detect_weight,
    inputs=gr.Image(type="pil", label="Upload or Capture Image"),
    outputs=[
        gr.Image(type="pil", label="Snapshot"),
        gr.Markdown(label="Result")
    ],
    title="⚖️ Auto Weight Logger",
    description="Upload or capture a digital scale snapshot to detect weight automatically.",
    allow_flagging="never"
)

if __name__ == "__main__":
    demo.launch()