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

from pydantic import SecretStr
from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage, AIMessage, ToolMessage
from langchain.tools import tool
from deepagents import create_deep_agent
from deepagents.backends import CompositeBackend, LocalShellBackend, FilesystemBackend

# DESIGN DECISION: Use OpenRouter LLM as required by the technical constraints.
# NECESSITY: The course forbids local LM endpoints and mandates OpenRouter for all LLM calls.
# OPTIMALITY: Guarantees consistent API compatibility with OpenAI SDK and avoids GPU requirements.
# ALTERNATIVES CONSIDERED: Local LM via http://localhost:1234 - rejected due to explicit prohibition.
llm = ChatOpenAI(
    model="openai/gpt-oss-20b:free",
    base_url="https://openrouter.ai/api/v1",
    api_key=os.getenv("OPENAI_API_KEY"),
    temperature=0.7,
)

# Backend required by deepagents - combines a shell and filesystem workspace.
backend = CompositeBackend(
    [
        LocalShellBackend(workspace_dir="./workspace"),
        FilesystemBackend(),
    ]
)

@tool
def get_price(product: str, city: str) -> str:
    """
    Получить примерную цену продукта в указанном городе.
    Возвращает таблицу в markdown-формате:
    | Продукт | Цена (руб.) | Магазин |
    """
    # DESIGN DECISION: Sub-agent is created inside the tool using the same LLM.
    # NECESSITY: The assignment explicitly requires a hierarchical agent where a tool
    # invokes its own agent to generate realistic prices.
    # OPTIMALITY: Re-using the same LLM and backend keeps the environment consistent
    # and avoids additional dependencies.
    # ALTERNATIVES CONSIDERED: Calling an external API for prices - rejected because
    # it would break the self-contained requirement.
    sub_agent = create_deep_agent(
        model=llm,
        tools=[],  # No further tools needed for price generation
        backend=backend,
        system_prompt=(
            "Ты суб-агент, который генерирует реалистичную цену продукта в заданном городе. "
            "Ответ дай в виде markdown-таблицы с колонками: Продукт, Цена (руб.), Магазин."
        ),
    )

    # Формируем запрос к суб-агенту
    query = f"Сгенерируй цену для продукта '{product}' в городе '{city}'."
    # Асинхронный вызов суб-агента
    async def invoke_sub() -> Dict[str, Any]:
        return await sub_agent.ainvoke(
            {"messages": [HumanMessage(content=query)]},
            {"configurable": {"thread_id": f"price-{product}-{city}"}},
        )

    # Запускаем цикл событий, если уже внутри async контекста
    try:
        loop = asyncio.get_running_loop()
        result = loop.create_task(invoke_sub())
        sub_result = asyncio.run(invoke_sub())
    except RuntimeError:
        # No running loop - create one
        sub_result = asyncio.run(invoke_sub())

    # Последнее сообщение суб-агента содержит таблицу
    price_table = sub_result["messages"][-1].content
    return price_table

# Главный агент
agent = create_deep_agent(
    model=llm,
    tools=[get_price],
    backend=backend,
    system_prompt="Ты помощник по планированию покупок.",
)

def format_message(msg: Any) -> str:
    """Привести сообщение к читаемому виду."""
    if isinstance(msg, AIMessage) or isinstance(msg, HumanMessage):
        return f"{msg.type.upper()}: {msg.content}"
    if isinstance(msg, ToolMessage):
        return f"TOOL CALL: {msg.name}({msg.args}) -> {msg.content}"
    # Fallback
    return str(msg)

async def main() -> None:
    user_query = "Помоги составить список покупок: молоко, хлеб, яблоки. Я нахожусь в Казани."
    result = await agent.ainvoke(
        {"messages": [HumanMessage(content=user_query)]},
        {"configurable": {"thread_id": "session-1"}},
    )

    # Вывод всей цепочки сообщений
    for i, message in enumerate(result["messages"]):
        print(f"--- Message {i + 1} ---")
        print(format_message(message))
        print()

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