app.py

import streamlit as st
import os
import time
import re
import requests
import json
from PIL import Image
from io import BytesIO
from urllib.parse import quote
from openai import OpenAI

# ------------------ App Configuration ------------------
st.set_page_config(page_title="Schlaeger Forrestdale TechDocAIA", layout="wide", initial_sidebar_state="collapsed")
st.title("📄 Schlaeger Forrestdale Document Assistant")
st.caption("Explore City of Armadale construction documents using AI + OCR 🧐")

# ------------------ Load API Key and Assistant ID ------------------
OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY")
ASSISTANT_ID = os.environ.get("ASSISTANT_ID")

if not OPENAI_API_KEY or not ASSISTANT_ID:
    st.error("❌ Missing secrets. Please set both OPENAI_API_KEY and ASSISTANT_ID in Hugging Face Space secrets.")
    st.stop()

client = OpenAI(api_key=OPENAI_API_KEY)

# ------------------ Session State Initialization ------------------
if "messages" not in st.session_state:
    st.session_state.messages = []
if "thread_id" not in st.session_state:
    st.session_state.thread_id = None
if "pending_prompt" not in st.session_state:
    st.session_state.pending_prompt = None
if "results" not in st.session_state:
    st.session_state.results = []
if "lightbox_url" not in st.session_state:
    st.session_state.lightbox_url = None

# ------------------ Sidebar ------------------
st.sidebar.header("ℹ️ Information")
if st.sidebar.button("🧹 Clear Chat"):
    st.session_state.messages = []
    st.session_state.thread_id = None
    st.session_state.pending_prompt = None
    st.session_state.results = []
    st.session_state.lightbox_url = None
    st.rerun()

# ------------------ Chat Input ------------------
user_prompt = st.chat_input("Ask about plans, drawings or components")
if user_prompt:
    st.session_state.messages.append({"role": "user", "content": user_prompt})

# ------------------ Assistant Query ------------------
if st.session_state.messages and st.session_state.messages[-1]["role"] == "user":
    try:
        if st.session_state.thread_id is None:
            thread = client.beta.threads.create()
            st.session_state.thread_id = thread.id

        client.beta.threads.messages.create(
            thread_id=st.session_state.thread_id,
            role="user",
            content=st.session_state.messages[-1]["content"]
        )

        run = client.beta.threads.runs.create(
            thread_id=st.session_state.thread_id,
            assistant_id=ASSISTANT_ID
        )

        with st.spinner("🤖 Parsing and responding..."):
            while True:
                run_status = client.beta.threads.runs.retrieve(
                    thread_id=st.session_state.thread_id,
                    run_id=run.id
                )
                if run_status.status in ("completed", "failed", "cancelled"):
                    break
                time.sleep(1)

        if run_status.status != "completed":
            st.error(f"⚠️ Assistant failed: {run_status.status}")
        else:
            messages = client.beta.threads.messages.list(thread_id=st.session_state.thread_id)
            for message in reversed(messages.data):
                if message.role == "assistant":
                    assistant_reply = message.content[0].text.value
                    st.session_state.messages.append({"role": "assistant", "content": assistant_reply})
                    try:
                        json_data = json.loads(assistant_reply.strip("`json "))
                        st.session_state.results = json_data
                    except:
                        st.session_state.results = []
                    break
        st.rerun()
    except Exception as e:
        st.error(f"❌ Error: {e}")

# ------------------ Sorting and Filters ------------------
if st.session_state.results:
    disciplines = sorted(set(d.get("discipline", "") for d in st.session_state.results))
    selected_discipline = st.selectbox("🌍 Filter by discipline", ["All"] + disciplines)
    page_size = 8
    page_num = st.number_input("Page", min_value=1, step=1, value=1)

    filtered_results = [r for r in st.session_state.results if selected_discipline == "All" or r.get("discipline") == selected_discipline]
    paged = filtered_results[(page_num - 1) * page_size : page_num * page_size]

    st.markdown("---")
    st.subheader("📂 Drawing Results")
    cols = st.columns(4)

    for i, item in enumerate(paged):
        with cols[i % 4]:
            st.markdown(f"**{item['drawing_number']}**")
            st.markdown(f"_Discipline: {item['discipline']}_")
            st.caption(item.get("summary", ""))
            for url in item.get("images", [])[:1]:
                if st.button("🖼️ View Image", key=f"view_{i}"):
                    st.session_state.lightbox_url = url

    if st.session_state.lightbox_url:
        st.markdown("---")
        st.image(st.session_state.lightbox_url, use_column_width=True, caption="🔍 Enlarged Preview")
        if st.button("❌ Close Viewer"):
            st.session_state.lightbox_url = None
            st.rerun()
else:
    for msg in st.session_state.messages:
        with st.chat_message(msg["role"]):
            st.markdown(msg["content"], unsafe_allow_html=True)