# solution.py

"""
Основной скрипт агента с памятью и подтверждением действий.
"""

import os
from typing import Any, Dict, Iterable, Tuple

from langchain_openai import ChatOpenAI
from langgraph.checkpoint.memory import MemorySaver
from langgraph.graph import StateGraph
from langgraph.prebuilt import create_react_agent
from rich.console import Console

# --------------------------------------------------------------------------- #
# 1. Конфигурация LLM и инструмента
# --------------------------------------------------------------------------- #

# Задаём ключ API OpenAI (можно задать переменной окружения)
os.environ.setdefault("OPENAI_API_KEY", "sk-...")

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

# Пример простого инструмента: поиск в интернете
def search(query: str) -> str:
    """Простейший инструмент поиска."""
    return f"Результаты поиска по запросу '{query}': ..."

search_tool = {
    "name": "search",
    "description": "Используется для поиска информации в интернете.",
    "func": search,
}

# --------------------------------------------------------------------------- #
# 2. Создаём агент с памятью и возможностью паузы перед вызовом инструмента
# --------------------------------------------------------------------------- #

memory = MemorySaver()

agent = create_react_agent(
    model=llm,
    tools=[search_tool],
    prompt="You are a helpful assistant.",
    checkpointer=memory,          # подключаем память
    interrupt_before=["tools"],   # ставим паузу перед каждым вызовом инструмента
)

# --------------------------------------------------------------------------- #
# 3. Вспомогательные функции
# --------------------------------------------------------------------------- #

console = Console()

def ask_and_run(user_input: Dict[str, Any], config: Dict[str, Any]) -> None:
    """
    Обрабатывает поток от агента, выводит токены и обрабатывает паузы.
    """
    # Переменная для хранения последнего состояния (для возобновления)
    last_state = None

    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["content"], end="")

        elif chunk_type == "updates":
            # Выводим информацию о вызове инструмента (если есть)
            for update in chunk_data:
                if isinstance(update, dict) and "tool_calls" in update.get("values", {}):
                    tool_call = update["values"]["messages"][-1].tool_calls[0]
                    console.print(f"\n[bold cyan]Инструмент:[/bold cyan] {tool_call['name']}({tool_call['args']})")

        # Проверяем, пришёл ли сигнал о паузе
        if "__interrupt__" in chunk_data and agent.get_state(config).next == ("tools",):
            # Получаем информацию о том, какой инструмент будет вызван
            state = agent.get_state(config)
            tool_call = state.values["messages"][-1].tool_calls[0]
            console.print(f"\n[bold magenta]Пауза:[/bold magenta] Агент хочет вызвать инструмент "
                          f"[green]{tool_call['name']}[/green] с аргументами {tool_call['args']}")
            # Спрашиваем подтверждение у пользователя
            answer = input("Разрешить? (Y/n): ").strip().lower()
            if answer in ("", "y", "yes"):
                console.print("[bold green]Продолжаем...[/bold green]")
                # Возобновляем работу агента с того места, где остановились
                ask_and_run(None, config)
                return
            else:
                console.print("[bold red]Отменено. Завершаем разговор.[/bold red]")
                break

        last_state = chunk_data

    console.print("\n[bold blue]Финальный ответ получен.[/bold blue]\n")

# --------------------------------------------------------------------------- #
# 4. Основной цикл чата
# --------------------------------------------------------------------------- #

def main() -> None:
    """
    Запускает бесконечный чат с агентом.
    """
    thread_id = "thread-1"
    config = {"configurable": {"thread_id": thread_id}}

    console.print("[bold underline]Запуск агента. Введите 'exit' для завершения.[/bold underline]")
    while True:
        user_input = input("\nВы: ")
        if user_input.lower() == "exit":
            console.print("[bold red]Завершение работы.[/bold red]")
            break

        # Формируем сообщение в формате, ожидаемом агентом
        message = {"messages": [{"role": "human", "content": user_input}]}
        ask_and_run(message, config)


if __name__ == "__main__":
    main()