import os
import asyncio
from typing import Optional, Dict, Any

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

# Инициализация консоли rich
console = Console()

# Инициализация LLM через OpenRouter
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,
)

# Backend для deepagents (необязательно, но удобно)
backend = CompositeBackend([
    LocalShellBackend(workspace_dir="./workspace"),
    FilesystemBackend(),
])

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

# Память разговора
memory = MemorySaver()

# Создание агента с памятью и паузой перед инструментом
agent = create_agent(
    model=llm,
    tools=[get_price],
    backend=backend,
    system_prompt="You are a helpful agent.",
    checkpointer=memory,
    interrupt_before=["tools"],
)

# Конфигурация разговора
config: Dict[str, Any] = {"configurable": {"thread_id": "conversation-1"}}

async def ask_and_run(user_input: Optional[Dict[str, Any]], config: Dict[str, Any]) -> None:
    """
    Запускает потоковое взаимодействие с агентом.
    Если агент останавливается перед вызовом инструмента, запрашивает подтверждение у пользователя.
    """
    async for chunk_type, chunk_data in agent.stream(
        user_input,
        config=config,
        stream_mode=["messages", "updates"],
    ):
        # Вывод токенов ответа
        if chunk_type == "messages":
            content = chunk_data.get("content", "")
            console.print(content, end="")

        # Вывод информации о вызове инструмента
        elif chunk_type == "updates":
            console.print(chunk_data)

        # Обнаружение паузы перед инструментом
        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("\n")
            console.print(f"{tool_call['name']}({tool_call['args']})")
            console.print("Агент хочет вызвать утилиту")
            answer = input("Разрешить? (Y/n): ")
            if answer.lower().strip() == "y":
                await ask_and_run(None, config)
                return
            else:
                console.print("Отменено")
                return

def main() -> None:
    console.print("\n--- --- ---\n")
    while True:
        user_input = input("\nВы: ")
        if user_input.lower().strip() == "exit":
            break
        # Запускаем асинхронную функцию
        asyncio.run(
            ask_and_run(
                {"messages": [{"role": "human", "content": user_input}]},
                config,
            )
        )
        console.print("\n--- --- ---\n")

if __name__ == "__main__":
    main()