# solution.py
"""
Основной скрипт агента с потоковым выводом.
Используются только библиотеки, указанные в условии:
    - langchain.agents.create_agent
    - langchain_ollama.ChatOllama
    - langchain.tools.tool
    - rich.print

Потоковый вывод реализован через agent.stream().
"""

# ────────────────────── подключение к LLM ──────────────────────
from langchain_ollama import ChatOllama
from langchain.agents import create_agent
from langchain.tools import tool
from rich import print as rprint

# Создаём объект модели Ollama (используем локальный сервер)
llm = ChatOllama(model="llama3.1")  # можно заменить на любой доступный модельный тег

# ────────────────────── определение инструмента с субагентом ──────────────────────
@tool
def get_price(product: str, city: str) -> str:
    """
    Симуляция запроса цены продукта в указанном городе.
    В реальном проекте здесь могла бы быть интеграция с API магазина.
    """
    # простая фиктивная логика для демонстрации работы инструмента
    prices = {
        ("молоко", "Казань"): 89,
        ("хлеб", "Казань"): 35,
        ("яблоки", "Казань"): 120,
    }
    price = prices.get((product.lower(), city.lower()), None)
    if price is None:
        return f"Цена на {product} в городе {city} не найдена."
    # форматируем таблицу
    table = (
        "| Продукт | Цена (руб.) | Магазин |\n"
        "|---------|-------------|---------|\n"
        f"| {product.title()}  | {price:<11} | Магнит   |\n"
    )
    return table

# ────────────────────── создание главного агента ──────────────────────
agent = create_agent(
    llm=llm,
    tools=[get_price],
    verbose=False,          # отключаем внутренний вывод LangChain
)

# ────────────────────── потоковый вывод ──────────────────────
def format_message(message) -> str:
    """
    Форматируем сообщение для вывода в консоль.
    Если у сообщения есть content – возвращаем его,
    иначе формируем строку вызова инструмента.
    """
    if message.content:
        return message.content
    # если нет content, но есть tool_calls – выводим их как строку
    if message.tool_calls:
        call = message.tool_calls[0]
        name = call["name"]
        args = call.get("args", {})
        return f"{name}({args})"
    return ""

def format_chunk_message(chunk):
    """
    Обрабатываем чанк типа 'messages'.
    Выводим токен сразу, добавляя разделитель при смене шага.
    """
    message, meta = chunk
    # Проверяем шаг агента
    if meta.get("langgraph_step") != format_chunk_message.step:
        format_chunk_message.step = meta["langgraph_step"]
        rprint("\n --- --- --- \n")
    content = message.content or ""
    if content:
        rprint(content, end="", flush=True)

# Инициализируем переменную шага
format_chunk_message.step = 0

# Запускаем потоковый вывод
stream = agent.stream(
    {
        "messages": [
            {"role": "human", "content": "Сколько стоит молоко в Казани?"}
        ]
    },
    stream_mode=["messages", "updates"],
)

for chunk_type, chunk_data in stream:
    if chunk_type == "messages":
        format_chunk_message(chunk_data)
    elif chunk_type == "updates":
        # Обрабатываем события обновления состояния (например, завершение шага модели)
        model_info = chunk_data.get("model")
        if model_info:
            last_msg = model_info["messages"][-1]
            rprint("\n" + format_message(last_msg) + "\n")

# Завершаем вывод
rprint("\n--- Конец ответа ---\n")