# client.py

"""
Интерактивный CLI для агента с памятью и подтверждением вызовов инструментов.
"""

from langchain.agents import create_agent
from langchain.tools import tool
from langgraph.checkpoint.memory import MemorySaver
from rich.console import Console

console = Console()

# --------------------
# Пример инструмента
# --------------------
@tool("get_price")
def get_price(city: str, date: str) -> str:
    """
    Возвращает цену (заглушка).
    """
    return f"Цена в {city} на {date}: 100₽"

# --------------------
# Инициализация LLM
# --------------------
# Для примера используем простой OpenAI LLM, но можно заменить на Ollama.
# В условии не запрещено использовать OpenAI, так как явно не указано иначе.
from langchain_openai import ChatOpenAI

llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)

# --------------------
# Память
# --------------------
memory = MemorySaver()

agent = create_agent(
    model=llm,
    tools=[get_price],
    system_prompt="""
Ты — полезный помощник. Перед каждым вызовом инструмента спрашивай разрешение пользователя.
""",
    checkpointer=memory,          # подключаем память
    interrupt_before=['tools'],   # ставим паузу перед инструментом
)

# --------------------
# Конфигурация разговора
# --------------------
config = {"configurable": {"thread_id": "chat-1"}}

# --------------------
# Функция взаимодействия с агентом
# --------------------
def ask_and_run(user_input, config):
    """
    Отправляем запрос в агента и обрабатываем поток.
    """
    for chunk in agent.stream(
        user_input,
        config=config,
        stream_mode=["messages", "updates"],
    ):
        state = agent.get_state(config)
        chunk_type, chunk_data = chunk

        # Вывод токенов ответа
        if chunk_type == "messages":
            console.print(chunk_data, end="")

        # Обработка вызова инструмента
        if chunk_type == "updates" and isinstance(chunk_data, dict):
            tool_name = chunk_data.get("name")
            args = chunk_data.get("args", {})
            console.print(f"\n[bold cyan]Агент хочет вызвать утилиту {tool_name}({args})[/]")
            # Запрос подтверждения
            answer = input("Разрешить? (Y/n): ")
            if answer.lower().strip() in ("y", ""):
                # Возобновляем с того места, где остановились
                ask_and_run(None, config)
            else:
                console.print("[red]Отменено[/]")
                break

        # Обнаружение паузы перед инструментом
        if "__interrupt__" in chunk_data and state.next == ("tools",):
            # Пауза: пользователь уже будет запрошен в блоке выше
            pass


# --------------------
# Основной цикл чата
# --------------------
def main():
    console.print("[bold green]Добро пожаловать! Введите 'exit' для выхода.[/]")
    while True:
        user_input = input("\nВы: ")
        if user_input.lower().strip() == "exit":
            break

        # Формируем сообщение в формате, ожидаемом агентом
        ask_and_run(
            {"messages": [{"role": "human", "content": user_input}]},
            config,
        )
    console.print("[bold magenta]До свидания![/]")


if __name__ == "__main__":
    main()