import chainlit as cl

from graph import workflow
from database import init_storage
import database

_graph = None

@cl.on_chat_start
async def start():
    global _graph
    try:
        await init_storage()
        if _graph is None:
            _graph = workflow.compile(checkpointer=database.checkpointer)
    except Exception as e:
        print(f"Ошибка при создании таблиц памяти: {e}")
    
    cl.user_session.set("thread_id", cl.user_session.get("id"))
    await cl.Message(content="Система запущена.").send()

@cl.on_message
async def main(message: cl.Message):
    thread_id = cl.user_session.get("thread_id")
    config = {"configurable": {"thread_id": thread_id}}
    
    # Запускаем LangGraph
    async for event in _graph.astream(
        {"messages": [("user", message.content)]}, 
        config=config
    ):
        pass

    final_state = await _graph.aget_state(config)
    last_msg = final_state.values["messages"][-1].content
    await cl.Message(content=last_msg).send()