import gradio as gr
import os
import json
from openai import OpenAI

# ========== Config ==========
OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY")
ASSISTANT_ID = os.environ.get("ASSISTANT_ID")
client = OpenAI(api_key=OPENAI_API_KEY)

# ========== Functions ==========
def query_assistant(message):
    try:
        thread = client.beta.threads.create()
        client.beta.threads.messages.create(
            thread_id=thread.id,
            role="user",
            content=message
        )
        run = client.beta.threads.runs.create(
            thread_id=thread.id,
            assistant_id=ASSISTANT_ID
        )

        # Poll run status
        while True:
            run_status = client.beta.threads.runs.retrieve(
                thread_id=thread.id,
                run_id=run.id
            )
            if run_status.status in ("completed", "failed", "cancelled"):
                break

        if run_status.status != "completed":
            return f"❌ Assistant failed: {run_status.status}", []

        messages = client.beta.threads.messages.list(thread_id=thread.id)
        for m in reversed(messages.data):
            if m.role == "assistant":
                try:
                    parsed = json.loads(m.content[0].text.value)
                    return None, parsed.get("results", [])
                except Exception as e:
                    return f"⚠️ Failed to parse assistant response: {e}", []
        return "⚠️ No assistant response found", []
    except Exception as e:
        return f"❌ Error: {e}", []

def render_cards(results):
    cards = []
    for r in results:
        title = r.get("drawing_title", "Untitled")
        summary = r.get("summary", "")
        pages = r.get("pages", [])
        image_elems = []
        for page in pages:
            url = page.get("public_image_url")
            label = f"{title} – Page {page.get('page_number')}"
            if url:
                image_elems.append(gr.Image(value=url, label=label, show_label=True))

        with gr.Column(scale=1):
            cards.append(
                gr.Group(
                    [
                        gr.Markdown(f"**{title}**\n\n**Summary:** {summary}"),
                        gr.Accordion("📂 View Drawing Pages", open=False, children=image_elems)
                    ]
                )
            )
    return cards

def on_query_submit(prompt):
    status, results = query_assistant(prompt)
    if status:
        return status, []
    return "✅ Found matching drawings", render_cards(results)

# ========== Interface ==========
with gr.Blocks(theme=gr.themes.Base(), title="Forrestdale Technical Drawing Assistant") as app:
    gr.Markdown("""
    # 🏗️ Forrestdale Technical Drawing Assistant
    Ask about plans, drawings or components (e.g. _Show me all electrical plans_)
    """)

    with gr.Row():
        query_input = gr.Textbox(placeholder="e.g. Show me all civil drainage plans", scale=5)
        query_button = gr.Button("🔍 Search", scale=1)

    status_output = gr.Markdown("", visible=True)
    result_display = gr.Column()

    query_button.click(fn=on_query_submit, inputs=query_input, outputs=[status_output, result_display])

# ========== Launch ==========
if __name__ == "__main__":
    app.launch()