"""Практическое задание №3: память разговора + подтверждение вызовов инструментов."""
from __future__ import annotations

import os

from dotenv import load_dotenv
from langchain.agents import create_agent
from langchain.tools import tool
from langchain_openai import ChatOpenAI
from langgraph.checkpoint.memory import MemorySaver
from rich.console import Console

load_dotenv()

console = Console()

llm = ChatOpenAI(
    model=os.getenv("OPENAI_MODEL", "openai/gpt-oss-20b:free"),
    base_url=os.getenv("OPENAI_BASE_URL", "https://openrouter.ai/api/v1"),
    api_key=os.getenv("OPENAI_API_KEY", "fake"),
    temperature=0.0,
)


@tool
def get_price(city: str, date: str = "сегодня") -> str:
    """Узнать примерную цену/стоимость покупок или погодные условия в городе на дату."""
    return f"{city}, {date}: ориентировочно 150–300 руб. за базовую корзину."


memory = MemorySaver()

agent = create_agent(
    model=llm,
    tools=[get_price],
    system_prompt="Ты помощник по планированию покупок и погоде. Помни контекст разговора.",
    checkpointer=memory,
    interrupt_before=["tools"],
)


def _print_pending_tool(config: dict) -> None:
    state = agent.get_state(config)
    messages = state.values.get("messages", [])
    if not messages:
        return
    last = messages[-1]
    tool_calls = getattr(last, "tool_calls", None) or []
    if not tool_calls:
        return
    tc = tool_calls[0]
    name = tc.get("name") if isinstance(tc, dict) else getattr(tc, "name", "?")
    args = tc.get("args") if isinstance(tc, dict) else getattr(tc, "args", {})
    console.print("\n --- --- --- ")
    console.print(f"{name}({args})")
    console.print(f"Агент хочет вызвать утилиту {name}({args})")


def _print_updates(chunk_data: dict) -> None:
    for node_name, update in chunk_data.items():
        if node_name == "__interrupt__":
            continue
        if not isinstance(update, dict):
            continue
        messages = update.get("messages", [])
        for msg in messages:
            tool_calls = getattr(msg, "tool_calls", None) or []
            for tc in tool_calls:
                name = tc.get("name") if isinstance(tc, dict) else getattr(tc, "name", "?")
                args = tc.get("args") if isinstance(tc, dict) else getattr(tc, "args", {})
                console.print(f"\n --- --- --- ")
                console.print(f"{name}({args})")


def ask_and_run(user_input: dict | None, config: dict) -> None:
    """Запуск или возобновление агента с обработкой паузы перед tools."""
    for chunk in agent.stream(
        user_input,
        config=config,
        stream_mode=["messages", "updates"],
    ):
        state = agent.get_state(config)

        if not isinstance(chunk, tuple) or len(chunk) != 2:
            continue

        chunk_type, chunk_data = chunk

        if chunk_type == "messages":
            if isinstance(chunk_data, tuple) and len(chunk_data) >= 1:
                token = chunk_data[0]
                content = getattr(token, "content", None)
                if content:
                    console.print(content, end="")

        if chunk_type == "updates" and isinstance(chunk_data, dict):
            _print_updates(chunk_data)

            if "__interrupt__" in chunk_data and state.next == ("tools",):
                _print_pending_tool(config)
                answer = input("Разрешить? (Y/n): ").strip().lower()
                if answer in ("", "y", "yes", "д", "да"):
                    ask_and_run(None, config)
                else:
                    console.print("Отменено")
                return

    console.print()


def main() -> None:
    config = {"configurable": {"thread_id": "разговор-1"}}
    console.print(
        "Чат с агентом (память + подтверждение tools). "
        "Введите 'exit' для выхода.\n"
    )

    while True:
        user_text = input("\nВы: ").strip()
        if user_text.lower() == "exit":
            break
        if not user_text:
            continue

        ask_and_run({"messages": [{"role": "human", "content": user_text}]}, config)


if __name__ == "__main__":
    main()