# solution.py

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

from langchain.chat_models import ChatOpenAI
from langchain.tools import BaseTool
from langgraph.checkpoint.memory import MemorySaver
from langgraph.graph import StateGraph
from langgraph.prebuilt import create_agent
from rich.console import Console

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

# Для работы нужен ключ OpenAI. Если его нет – создаём фиктивный LLM.
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
if OPENAI_API_KEY:
    llm = ChatOpenAI(temperature=0, openai_api_key=OPENAI_API_KEY)
else:
    # Минимальный заглушка‑LLM (не делает ничего полезного)
    class DummyLLM:
        def __call__(self, *args, **kwargs):
            return {"content": "Dummy response"}
    llm = DummyLLM()

# Пример простого инструмента
class EchoTool(BaseTool):
    name: str = "echo"
    description: str = "Возвращает переданный текст без изменений."
    def _run(self, text: str) -> str:
        return f"Echo: {text}"

my_tool = EchoTool()

# ------------------------------------------------------------
# Создание агента с памятью и паузой
# ------------------------------------------------------------

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

console = Console()

# ------------------------------------------------------------
# Функция взаимодействия с пользователем
# ------------------------------------------------------------

def ask_and_run(user_input: Optional[str], config: Dict[str, Any]) -> None:
    """
    Запускает потоковый диалог с агентом. Если агент хочет вызвать инструмент,
    пользователь получает запрос на подтверждение.
    """
    # Если user_input – None, продолжаем с того места, где остановились
    stream = agent.stream(
        {"messages": [] if user_input is None else [{"role": "user", "content": user_input}]},
        config=config,
        stream_mode=["messages", "updates"],
    )

    for chunk_type, chunk_data in stream:
        # 1. Печать токенов (части сообщения)
        if chunk_type == "messages":
            console.print(chunk_data["content"], end="", style="cyan")
            continue

        # 2. Обработка вызова инструмента
        if chunk_type == "updates" and "tool_calls" in chunk_data:
            for tool_call in chunk_data["tool_calls"]:
                name = tool_call["name"]
                args = tool_call.get("args", {})
                console.print(f"\n[bold magenta]Tool call:[/bold magenta] {name}({args})")
            continue

        # 3. Обнаружение паузы перед инструментом
        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]
            name = tool_call["name"]
            args = tool_call.get("args", {})
            console.print(f"\n[bold yellow]Пауза:[/bold yellow] {name}({args})")
            answer = input("\nРазрешить вызов инструмента? (Y/n): ").strip().lower()
            if answer in ("y", ""):
                # Продолжаем с того же состояния
                ask_and_run(None, config)
            else:
                console.print("[red]Вызов отменён.[/red]")
            break

# ------------------------------------------------------------
# Точка входа
# ------------------------------------------------------------

if __name__ == "__main__":
    thread_id = "thread-1"
    config = {"configurable": {"thread_id": thread_id}}
    console.print("[bold green]Запуск агента. Введите сообщение, а затем Enter.[/bold green]")
    while True:
        try:
            user_msg = input("\nВы: ")
            if user_msg.lower() in ("exit", "quit"):
                console.print("[bold red]Завершение работы.[/bold red]")
                break
            ask_and_run(user_msg, config)
        except KeyboardInterrupt:
            console.print("\n[red]Прервано пользователем. Выход.[/red]")
            break