import streamlit as st
import os
import json
import time
import requests
from openai import OpenAI

# ------------------ App Configuration ------------------
st.set_page_config(page_title="Forrestdale Drawing Assistant", layout="wide")
st.title("📐 Forrestdale Technical Drawing Assistant")

# ------------------ 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 OPENAI_API_KEY and ASSISTANT_ID in secrets.")
    st.stop()

client = OpenAI(api_key=OPENAI_API_KEY)

# ------------------ Session State Init ------------------
if "thread_id" not in st.session_state:
    st.session_state.thread_id = None
if "results" not in st.session_state:
    st.session_state.results = []

# ------------------ Query Bar ------------------
prompt = st.text_input("Ask about plans, drawings or components", placeholder="e.g. Show me all electrical plans")

if prompt:
    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=prompt
    )

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

    with st.spinner("Thinking..."):
        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":
        messages = client.beta.threads.messages.list(thread_id=st.session_state.thread_id)
        for message in reversed(messages.data):
            if message.role == "assistant":
                try:
                    json_data = json.loads(message.content[0].text.value)
                    st.session_state.results = json_data["results"]
                except Exception as e:
                    st.error("⚠️ Could not parse assistant response as JSON.")
                break
    else:
        st.error(f"⚠️ Assistant failed: {run_status.status}")

# ------------------ Display Results as Cards ------------------
if st.session_state.results:
    cols = st.columns(4)
    for i, result in enumerate(st.session_state.results):
        with cols[i % 4]:
            st.markdown("""
                <div style='border:1px solid #444; padding:10px; border-radius:10px; background:#111;'>
                    <h5 style='color:#fdd;'>📁 {drawing_number}</h5>
                    <b>Summary:</b> {summary}<br>
                    <img src="{image_url}" style="width:100%; margin-top:10px; border-radius:6px;">
                </div>
            """.format(
                drawing_number=result.get("drawing_number", "Untitled"),
                summary=result.get("summary", "No summary available."),
                image_url=result.get("image_url", "")
            ), unsafe_allow_html=True)