import spaces  # noqa: F401  — must be imported first, before torch/gradio pull it in transitively, per ZeroGPU requirements
import gradio as gr
from agents import build_graph, run_task

# Compiled once per process — the graph itself holds no per-user state,
# all mutable data lives in gr.State (per-session) so concurrent users
# on the same Space don't interfere with each other.
graph_app = build_graph()


def handle_message(message, history, state):
    state = state or {}
    try:
        result = run_task(graph_app, message, state)
    except Exception as e:
        result = f"Agent error: {e}"
    return result, state


with gr.Blocks(title="Multi-Agent Assistant") as demo:
    gr.Markdown("# Multi-Agent Assistant\nManager agent routes to search, code, and data specialists.")

    session_state = gr.State({})

    chatbot = gr.Chatbot(height=450)
    msg = gr.Textbox(placeholder="Ask something that needs research, code, or data analysis...")
    clear = gr.Button("Clear")

    def respond(message, chat_history, state):
        chat_history = chat_history or []
        chat_history.append({"role": "user", "content": message})
        reply, state = handle_message(message, chat_history, state)
        chat_history.append({"role": "assistant", "content": reply})
        return "", chat_history, state

    msg.submit(respond, [msg, chatbot, session_state], [msg, chatbot, session_state])
    clear.click(lambda: ([], {}), None, [chatbot, session_state])

if __name__ == "__main__":
    print("Starting Gradio server on 0.0.0.0:7860 ...", flush=True)
    demo.launch(server_name="0.0.0.0", server_port=7860)