# streamlit_app.py
import streamlit as st
import requests
import json
from PIL import Image
import io
import base64

API_URL = "http://localhost:7860/process-image"  # Change if hosted elsewhere

st.set_page_config(page_title="Flowchart to English", layout="wide")
st.title("📄 Flowchart to Plain English")

# Debug mode switch
debug_mode = st.toggle("🔧 Show Debug Info", value=False)

uploaded_file = st.file_uploader("Upload a flowchart image", type=["png", "jpg", "jpeg"])

if uploaded_file:
    # Resize image for smaller canvas
    image = Image.open(uploaded_file)
    max_width = 600
    ratio = max_width / float(image.size[0])
    new_height = int((float(image.size[1]) * float(ratio)))
    resized_image = image.resize((max_width, new_height))
    st.image(resized_image, caption="📤 Uploaded Image", use_container_width=False)

    if st.button("🔍 Analyze Flowchart"):
        progress = st.progress(0, text="Sending image to backend...")

        try:
            response = requests.post(
                API_URL,
                files={"file": uploaded_file.getvalue()},
                data={"debug": str(debug_mode).lower()}
            )
            progress.progress(50, text="Processing detection, OCR, and reasoning...")

            if response.status_code == 200:
                data = response.json()
                progress.progress(80, text="Generating explanation using LLM...")

                # Optional: Visualize bounding boxes
                if debug_mode and data.get("yolo_vis"):
                    st.markdown("### 🖼️ YOLO Debug Bounding Boxes")
                    vis_bytes = base64.b64decode(data["yolo_vis"])
                    vis_img = Image.open(io.BytesIO(vis_bytes))
                    st.image(vis_img, caption="YOLO Detected Boxes", use_container_width=True)

                # Optional: show logs
                if debug_mode and "debug" in data:
                    st.markdown("### 🧪 Debug Pipeline Info")
                    st.code(data["debug"], language="markdown")

                # Display results in 2 columns
                col1, col2 = st.columns(2)
                with col1:
                    st.subheader("🧠 Flowchart JSON")
                    st.json(data["flowchart"])
                with col2:
                    st.subheader("📝 English Summary")
                    st.markdown(data["summary"])

                progress.progress(100, text="Done!")
            else:
                st.error(f"Something went wrong: {response.status_code}")
        except Exception as e:
            st.error(f"An error occurred: {e}")
else:
    st.info("Upload a flowchart image to begin.")