import gradio as gr
import openai
import json
import os
import re

# Set your OpenAI API key
openai.api_key = os.getenv("OPENAI_API_KEY")

# OpenAI Assistant ID
ASSISTANT_ID = "asst_mDU549QCoU616nB69g5ZwI58"

# Load full product catalog
with open("cleaned_product_catalog.json", "r") as f:
    PRODUCT_CATALOG = {item["ProductCode"]: item for item in json.load(f)}

# Global session state
cart = []
thread_id = None

# 🧠 Smart JSON extractor that handles both ```json``` and loose JSON objects
def extract_cart_json_from_text(text):
    blocks = []

    # 1. Try to find ```json ... ``` blocks
    code_blocks = re.findall(r"```json(.*?)```", text, re.DOTALL)
    for block in code_blocks:
        try:
            blocks.append(json.loads(block.strip()))
        except Exception:
            continue

    # 2. Fall back to finding any JSON-like objects
    if not blocks:
        loose_jsons = re.findall(r"\{[\s\S]+?\}", text)
        for snippet in loose_jsons:
            try:
                candidate = json.loads(snippet)
                if "ProductCode" in candidate and "Description" in candidate:
                    blocks.append(candidate)
            except Exception:
                continue

    return blocks

# 🔄 Communicate with OpenAI Assistant
def send_to_assistant(message):
    global thread_id
    if thread_id is None:
        thread = openai.beta.threads.create()
        thread_id = thread.id

    openai.beta.threads.messages.create(
        thread_id=thread_id,
        role="user",
        content=message
    )

    run = openai.beta.threads.runs.create(
        thread_id=thread_id,
        assistant_id=ASSISTANT_ID
    )

    while True:
        run_status = openai.beta.threads.runs.retrieve(
            thread_id=thread_id,
            run_id=run.id
        )
        if run_status.status == "completed":
            break

    messages = openai.beta.threads.messages.list(thread_id=thread_id)
    return messages.data[0].content[0].text.value

# 📦 Handle chat + add to cart
def assistant_response(message, chat_history):
    global cart
    response = send_to_assistant(message)

    extracted_items = extract_cart_json_from_text(response)

    if extracted_items:
        for item in extracted_items:
            cart.append(item)

    chat_history.append((message, response))
    return "", chat_history, json.dumps(cart, indent=2)

# 🧼 Clear cart
def clear_cart():
    global cart
    cart.clear()
    return json.dumps(cart, indent=2)

# ✅ Finalise order
def submit_order():
    global cart
    if not cart:
        return "No items in cart, mate!"

    customer = cart[0].get("Customer", {}) if cart else {}
    discount_rate = customer.get("Discount", 0)

    # Safe subtotal calc
    subtotal = 0
    for item in cart:
        try:
            price = float(item["Price"])
        except:
            price = 0.0
        subtotal += price * item.get("Quantity", 1)

    discount_amt = round(subtotal * discount_rate, 2)
    total = round(subtotal - discount_amt, 2)

    # Save final order
    order_summary = {
        "Customer": customer,
        "Products": [
            {k: v for k, v in item.items() if k != "Customer"}
            for item in cart
        ],
        "Totals": {
            "Subtotal": round(subtotal, 2),
            "DiscountRate": discount_rate,
            "DiscountAmount": discount_amt,
            "TotalAfterDiscount": total
        }
    }

    with open("order_summary.json", "w") as f:
        json.dump(order_summary, f, indent=2)

    # Generate friendly summary
    response_text = f"✅ Sweet as, {customer.get('Name', 'mate')} — your order’s locked in!\n\nHere’s what’s in the kit:\n\n"
    for item in cart:
        response_text += f"- {item['ProductCode']} — {item['Description']} — 💰 ${item['Price']}\n"
    response_text += f"\n💲 **Subtotal:** ${subtotal:.2f}\n"
    response_text += f"🎯 **Discount ({int(discount_rate*100)}%):** -${discount_amt:.2f}\n"
    response_text += f"✅ **Total After Discount:** ${total:.2f}\n\n"
    response_text += "We’ll get this sorted faster than a roo in a paddock. You’ll hear from us soon!"

    cart.clear()

    return response_text

# 🖼️ Gradio Interface
with gr.Blocks() as demo:
    gr.Markdown("# 🧢 Westate Sales Assistant 🇦🇺\nTalk to your hose mate and build your order!")

    with gr.Row():
        with gr.Column(scale=3):
            chatbot = gr.Chatbot()
            msg = gr.Textbox(label="Type your question here", placeholder="Ask Willy anything...")
            send = gr.Button("Send")

        with gr.Column(scale=2):
            cart_display = gr.Textbox(label="🛒 Cart (Full JSON)", lines=20)
            clear_btn = gr.Button("Clear Cart")
            checkout_btn = gr.Button("Checkout")

    send.click(assistant_response, inputs=[msg, chatbot], outputs=[msg, chatbot, cart_display])
    clear_btn.click(fn=clear_cart, outputs=cart_display)
    checkout_btn.click(fn=submit_order, outputs=cart_display)

demo.launch()