"""Gradio chat UI for the Meridian support chatbot.

Run with:  uv run python -m app.ui
"""

import asyncio
import uuid
from contextlib import asynccontextmanager

import gradio as gr
from agents import SQLiteSession

from .agent import build_agent
from .chat import chat_turn
from .config import Config
from .mcp_client import build_mcp_server

CONFIG = Config()


@asynccontextmanager
async def lifespan():
    """Connect MCP once, share across all chat turns; clean up on exit."""
    server = build_mcp_server(CONFIG.mcp_server_url)
    await server.connect()
    try:
        yield server
    finally:
        await server.cleanup()


def build_ui() -> gr.Blocks:
    with gr.Blocks(title="Meridian Support") as demo:
        gr.Markdown("# Meridian Electronics — Support Chat")
        cost_badge = gr.Markdown("**Conversation cost:** $0.00000")

        chatbot = gr.Chatbot(height=480)
        msg = gr.Textbox(placeholder="Ask about products, orders, or your account...", show_label=False)
        clear = gr.Button("Reset conversation")

        # Per-browser-session state.
        agent_state = gr.State()
        session_state = gr.State()
        running_cost = gr.State(0.0)

        async def initialize():
            # Lazy: a single shared MCP server, attached to a fresh agent + session per browser.
            # For the prototype we accept that the MCP server reconnects per Gradio process restart.
            server = build_mcp_server(CONFIG.mcp_server_url)
            await server.connect()
            agent = build_agent(CONFIG.openai_model, server)
            session = SQLiteSession(f"meridian-{uuid.uuid4().hex[:8]}")
            return agent, session, 0.0, "**Conversation cost:** $0.00000"

        async def respond(user_message, history, agent, session, cost_so_far):
            history = (history or []) + [
                {"role": "user", "content": user_message},
                {"role": "assistant", "content": ""},
            ]
            cost_label = f"**Conversation cost:** ${cost_so_far:.5f}"
            yield history, cost_so_far, cost_label, ""

            text = ""
            async for partial, cost in chat_turn(user_message, agent, session):
                text = partial
                history[-1]["content"] = text
                if cost is None:
                    yield history, cost_so_far, cost_label, ""
                else:
                    cost_so_far = cost_so_far + cost.usd
                    cost_label = (
                        f"**Conversation cost:** ${cost_so_far:.5f}  "
                        f"·  last turn: {cost.format()}"
                    )
                    yield history, cost_so_far, cost_label, ""

        async def reset():
            session = SQLiteSession(f"meridian-{uuid.uuid4().hex[:8]}")
            return [], session, 0.0, "**Conversation cost:** $0.00000"

        demo.load(initialize, outputs=[agent_state, session_state, running_cost, cost_badge])
        msg.submit(
            respond,
            inputs=[msg, chatbot, agent_state, session_state, running_cost],
            outputs=[chatbot, running_cost, cost_badge, msg],
        )
        clear.click(reset, outputs=[chatbot, session_state, running_cost, cost_badge])

    return demo


demo = build_ui()


if __name__ == "__main__":
    demo.queue().launch()