import gradio as gr
from PIL import Image
from datetime import datetime
import pytz
import os
from simple_salesforce import Salesforce
from ocr_engine import extract_weight
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

# Salesforce connection
sf = Salesforce(
    username=os.getenv("SF_USERNAME"),
    password=os.getenv("SF_PASSWORD"),
    security_token=os.getenv("SF_TOKEN")
)

def process_image(image):
    if image is None:
        return "❌ No image provided", "", None, gr.update(visible=True)

    try:
        weight = extract_weight(image)
        ist = pytz.timezone('Asia/Kolkata')
        now = datetime.now(ist)
        timestamp_display = now.strftime("%Y-%m-%d %H:%M:%S IST")
        timestamp_iso = now.isoformat()

        if not weight or "No valid" in weight:
            return "❌ Unable to detect. Try again with a clearer image.", "", image, gr.update(visible=True)

        # Create Salesforce record
        sf.Weight_Log__c.create({
            'Captured_Weight__c': float(weight.replace("kg", "").strip()),
            'Captured_At__c': timestamp_iso,
            'Device_ID__c': 'Device-01',  # Static or dynamic device ID
            'Snapshot_Image__c': 'Manual Entry',
            'Status__c': 'Captured'
        })

        return weight, timestamp_display, image, gr.update(visible=False)

    except Exception as e:
        return f"Error: {str(e)}", "", None, gr.update(visible=True)

# Gradio UI
with gr.Blocks(css=".gr-button {background-color: #2e7d32 !important; color: white !important;}") as demo:
    gr.Markdown("""
    <h1 style='text-align: center; color: #2e7d32;'>📷 Auto Weight Logger</h1>
    <p style='text-align: center;'>Upload or capture a digital weight image. Detects weight using AI OCR and stores the result in Salesforce.</p>
    <hr style='border: 1px solid #ddd;'/>
    """)

    with gr.Row():
        image_input = gr.Image(type="pil", label="📁 Upload or Capture Image")

    detect_btn = gr.Button("🚀 Detect Weight")

    with gr.Row():
        weight_out = gr.Textbox(label="📦 Detected Weight", placeholder="e.g., 97.9 kg", show_copy_button=True)
        time_out = gr.Textbox(label="🕒 Captured At (IST)", placeholder="e.g., 2025-07-01 12:00:00")

    snapshot = gr.Image(label="📸 Snapshot Preview")
    retake_btn = gr.Button("🔁 Retake / Try Again", visible=False)

    detect_btn.click(fn=process_image, inputs=image_input, outputs=[weight_out, time_out, snapshot, retake_btn])
    retake_btn.click(fn=lambda: ("", "", None, gr.update(visible=False)),
                     inputs=[], outputs=[weight_out, time_out, snapshot, retake_btn])

demo.launch()