# main.py
# Реализация агента с памятью и подтверждением действий

import os
from rich.console import Console
from langgraph.checkpoint.memory import MemorySaver
from langgraph import create_agent
from langchain_ollama import ChatOllama
from langchain.tools import Tool

# Rich console for pretty output
console = Console()

# ---------- Tool definition ----------

def get_price(city: str, date: str) -> str:
    """Return a mock price for a given city and date."""
    return f"Price in {city} on {date} is $100"

price_tool = Tool(
    name="get_price",
    description="Retrieve the price for a specified city and date.",
    func=get_price,
)

# ---------- LLM and Agent setup ----------
llm = ChatOllama(model="llama3")
memory = MemorySaver()
agent = create_agent(
    model=llm,
    tools=[price_tool],
    system_prompt="You are a helpful assistant that can call the get_price tool.",
    checkpointer=memory,
    interrupt_before=["tools"],  # pause before each tool invocation
)

# ---------- Helper function to handle streaming and confirmation ----------

def ask_and_run(user_input, config):
    for chunk_type, chunk_data in agent.stream(
        user_input if user_input is not None else None,
        config=config,
        stream_mode=["messages", "updates"],
    ):
        if chunk_type == "messages":
            console.print(chunk_data["content"], end="")
            console.flush()
            continue
        if chunk_type == "updates":
            console.print(f"\n[bold cyan]Update:[/bold cyan] {chunk_data}")
            continue
        if "__interrupt__" in chunk_data and agent.get_state(config).next == ("tools",):
            state = agent.get_state(config)
            last_msg = state.values["messages"][-1]
            tool_call = last_msg.tool_calls[0]
            tool_name = tool_call["name"]
            tool_args = tool_call["args"]
            console.print(f"\n\n[bold yellow]Agent wants to call tool:[/bold yellow] {tool_name}({tool_args})")
            answer = input("Разрешить? (Y/n): ")
            if answer.lower().strip() == "y":
                ask_and_run(None, config)
            else:
                console.print("[red]Отменено[/red]")
            break

# ---------- Main chat loop ----------
if __name__ == "__main__":
    thread_id = os.getenv("THREAD_ID", "chat-1")
    config = {"configurable": {"thread_id": thread_id}}
    console.print("[bold green]Добро пожаловать в интерактивный агент![/bold green]")
    console.print("Введите 'exit' для выхода.")
    while True:
        user_input = input("\nВы: ")
        if user_input.lower().strip() == "exit":
            console.print("[bold magenta]До свидания![/bold magenta]")
            break
        message = {"messages": [{"role": "human", "content": user_input}]}
        ask_and_run(message, config)
        console.print("\n--- --- ---")