from langgraph import create_agent
from langgraph.checkpoint.memory import MemorySaver
from langchain_openai import ChatOpenAI
from rich.console import Console
from langgraph.tools import tool
import asyncio

console = Console()

@tool
def get_price(params: dict) -> str:
    """Return a mock price for a city and date."""
    city = params.get("city", "unknown")
    date = params.get("date", "unknown")
    return f"Цена в {city} на {date}: 1000₽"

llm = ChatOpenAI(
    base_url="http://localhost:11434/v1",
    api_key="ollama",
    model="llama3",
)

memory = MemorySaver()

agent = create_agent(
    model=llm,
    tools=[get_price],
    system_prompt="You are a helpful assistant.",
    checkpointer=memory,
    interrupt_before=["tools"],
)

config = {"configurable": {"thread_id": "conversation-1"}}

async def ask_and_run(user_input, config):
    try:
        async for chunk in agent.stream(user_input, config=config, stream_mode=["messages", "updates"]):
            chunk_type, chunk_data = chunk
            if chunk_type == "messages":
                console.print(chunk_data, end="")
            elif chunk_type == "updates":
                if "__interrupt__" in chunk_data and agent.get_state(config).next == ("tools",):
                    state = agent.get_state(config)
                    tool_calls = state.values.get("messages", [])[-1].tool_calls
                    if tool_calls:
                        last_call = tool_calls[0]
                        console.print(f"\n\nTool call detected: {last_call['name']}{last_call['args']}")
                        answer = input("Разрешить? (Y/n): ")
                        if answer.lower().strip() == "y":
                            await ask_and_run(None, config)
                        else:
                            console.print("Отменено")
                            break
    except Exception as e:
        console.print(f"[red]Ошибка во время выполнения:[/red] {e}")
        raise

def main():
    console.print("Добро пожаловать! Введите 'exit' для выхода.")
    while True:
        user_input = input("\nВы: ")
        if user_input.lower() == "exit":
            break
        msg = {"messages": [{"role": "human", "content": user_input}]}
        asyncio.run(ask_and_run(msg, config))
        console.print("\n---")

if __name__ == "__main__":
    main()