import os
import asyncio
from rich.console import Console
from rich.text import Text
from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage
from langchain.tools import tool
from deepagents import create_deep_agent
from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend
from langgraph.checkpoint.memory import MemorySaver

# ---------- Консоль ----------
console = Console()

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

# ---------- Инструменты ----------
@tool
def get_price(city: str, date: str) -> str:
    """
    Возвращает условную цену погоды для указанного города и даты.
    Параметры:
        city: название города
        date: дата в произвольном формате (например, "сегодня", "завтра")
    """
    # Здесь могла бы быть реальная логика обращения к API погоды.
    return f"Цена погоды в {city} на {date}: 42 условных единиц."

# ---------- Бэкенд ----------
backend = CompositeBackend(
    [
        LocalShellBackend(workspace_dir="./workspace"),
        FilesystemBackend(),
    ]
)

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

# ---------- Агент ----------
agent = create_deep_agent(
    model=llm,
    tools=[get_price],
    backend=backend,
    system_prompt="You are a helpful assistant. Use the provided tools when needed.",
    checkpointer=memory,
    interrupt_before=["tools"],  # пауза перед каждым вызовом инструмента
)

# ---------- Функция обработки одного запроса ----------
async def ask_and_run(user_input, config):
    """
    user_input: dict вида {"messages": [{"role": "human", "content": "..."}]}
                или None (для возобновления после подтверждения)
    config: dict с ключом "configurable" -> {"thread_id": "..."}
    """
    # Если передан None, просто продолжаем без новых сообщений
    input_payload = user_input if user_input is not None else {}

    async for chunk in agent.astream(
        input_payload,
        config,
        stream_mode=["messages", "updates"],
    ):
        # Получаем актуальное состояние после каждого чанка
        state = agent.get_state(config)

        chunk_type, chunk_data = chunk

        # 1. Вывод сообщений (поток токенов)
        if chunk_type == "messages":
            # chunk_data содержит список сообщений, но мы выводим только новые токены
            # Для простоты выводим полное содержание последнего сообщения, если оно от ассистента
            last_msg = chunk_data[-1]
            if getattr(last_msg, "role", None) == "assistant":
                console.print(Text(last_msg.content, style="green"))

        # 2. Вывод обновлений (инструменты)
        if chunk_type == "updates":
            # Здесь можно вывести информацию о вызове инструмента, если она есть
            pass  # в данном простом примере ничего не делаем

        # 3. Обработка прерывания перед инструментом
        if "__interrupt__" in chunk_data and state.next == ("tools",):
            # Получаем информацию о запланированном вызове инструмента
            tool_msg = state.values["messages"][-1]
            tool_call = tool_msg.tool_calls[0]
            tool_name = tool_call["name"]
            tool_args = tool_call["args"]
            console.print(f"[bold yellow]Agent wants to call {tool_name}{tool_args}[/]")

            answer = input("Разрешить? (Y/n): ").strip().lower()
            if answer == "n" or answer == "no":
                console.print("[red]Отменено[/]")
                # Прерываем текущий поток, не вызывая инструмент
                break
            else:
                # Возобновляем поток без нового пользовательского сообщения
                await ask_and_run(None, config)
                # После возврата продолжаем обработку (обычно поток уже завершён)
                return

# ---------- Основной цикл ----------
async def main():
    thread_id = "conversation-1"
    config = {"configurable": {"thread_id": thread_id}}

    console.print("[bold cyan]=== Чат с подтверждением инструментов ===[/]")
    while True:
        user_text = input("\nВы: ").strip()
        if user_text.lower() in {"exit", "quit"}:
            console.print("[bold magenta]Завершение работы.[/]")
            break

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

if __name__ == "__main__":
    asyncio.run(main())