"""
AI-агент для планирования списка покупок.

Иерархический агент на LangChain + локальная LLM (LM Studio).
- Главный агент принимает список продуктов
- Инструмент get_price использует субагент для получения цен
- Итоговая таблица с ценами и общей стоимостью
"""

import os
from langchain_core.tools import tool
from langchain_openai import ChatOpenAI
from langgraph.prebuilt import create_react_agent


# Подключение к LLM через OpenAI-совместимый API (LM Studio)
llm = ChatOpenAI(
    model=os.environ.get("LM_STUDIO_MODEL", "local-model"),
    openai_api_base=os.environ.get("LM_STUDIO_API_BASE", "http://localhost:1234/v1"),
    openai_api_key=os.environ.get("LM_STUDIO_API_KEY", "not-needed"),
    temperature=0,
)


@tool
def get_price(product: str, city: str) -> str:
    """Получить цену продукта в указанном городе. Возвращает таблицу с ценами.

    Args:
        product: Название продукта
        city: Город для поиска цен
    """
    # Субагент для поиска цен
    sub_agent = create_react_agent(
        model=llm,
        tools=[],
        prompt=(
            "Ты помощник по поиску цен на продукты. "
            "Сгенерируй реалистичную цену для указанного продукта в указанном городе. "
            "Верни результат строго в формате таблицы:\n"
            "| Продукт | Цена | Магазин |\n"
            "| {product} | {цена} руб | {магазин} |"
        ),
    )

    result = sub_agent.invoke({
        "messages": [{"role": "user", "content": f"Найди цену для: {product} в городе {city}"}]
    })

    # Извлекаем ответ из последнего сообщения
    last_message = result["messages"][-1]
    return last_message.content if hasattr(last_message, "content") else str(last_message)


def main():
    """Главный агент — помощник по планированию покупок."""

    # Главный агент с инструментом get_price
    agent = create_react_agent(
        model=llm,
        tools=[get_price],
        prompt="Ты помощник по планированию покупок. Помоги составить список покупок с ценами.",
    )

    # Запрос пользователя
    user_request = "Помоги составить список покупок: молоко, хлеб, яблоки. Я нахожусь в Казани."

    print("=" * 60)
    print("  Запрос:", user_request)
    print("=" * 60)

    # Запуск агента с выводом всех промежуточных сообщений
    result = agent.invoke(
        {"messages": [{"role": "user", "content": user_request}]},
    )

    # Вывод всех сообщений (промежуточные вызовы + финальный ответ)
    print("\n--- Все сообщения ---\n")
    for i, msg in enumerate(result["messages"]):
        role = msg.__class__.__name__
        content = msg.content if hasattr(msg, "content") else str(msg)
        print(f"[{i}] {role}:")
        print(f"    {content}")
        print()

    # Финальный ответ
    final = result["messages"][-1]
    print("=" * 60)
    print("  Финальный ответ:")
    print("=" * 60)
    print(final.content if hasattr(final, "content") else str(final))


if __name__ == "__main__":
    main()