import streamlit as st
from PIL import Image, UnidentifiedImageError
import io
from datetime import datetime
import pytz
from ocr_engine import extract_weight_from_image

st.set_page_config(page_title="⚖️ Auto Weight Logger", layout="centered")
st.title("⚖️ Auto Weight Logger")

if "camera_key" not in st.session_state:
    st.session_state["camera_key"] = 0

# UI: Upload or Capture Image
col1, col2 = st.columns(2)
with col1:
    uploaded_file = st.file_uploader("Upload Image", type=["jpg", "jpeg", "png"])
with col2:
    if st.button("📸 Retake Photo"):
        st.session_state["camera_key"] += 1
    camera_image = st.camera_input("Capture Image", key=st.session_state["camera_key"])

# Load image
image = None
if uploaded_file:
    try:
        image = Image.open(uploaded_file)
    except UnidentifiedImageError:
        st.error("Invalid image format.")
elif camera_image:
    try:
        image = Image.open(camera_image)
    except UnidentifiedImageError:
        st.error("Invalid camera image.")

# Process image
if image:
    with st.spinner("Detecting weight..."):
        weight, confidence = extract_weight_from_image(image)

    # Timestamp
    ist_time = datetime.now(pytz.timezone("Asia/Kolkata")).strftime("%Y-%m-%d %H:%M:%S")

    # Output
    st.markdown("### 📅 Captured At (IST)")
    st.info(ist_time)

    st.markdown("### 🖼️ Snapshot")
    st.image(image, width=400)

    st.markdown("### ⚖️ Detected Weight")
    if confidence > 0:
        st.success(f"Detected Weight: **{weight}**  \nConfidence: `{confidence:.2f}`")
    else:
        st.error("No weight detected. Try with a clearer image.")

    # (Optional) Send to Salesforce button
    st.markdown("🔗 [Send to Salesforce](#)", unsafe_allow_html=True)