# solution.py

"""
Основной скрипт агента с потоковым выводом.
Используется стек: LangChain, OpenAI и rich.
"""

from langchain import hub
from langchain.agents import initialize_agent, AgentType
from langchain.tools import tool
from langchain_community.llms.openai import OpenAI
from rich import print as rprint

# ---------- 1. Подключаем LLM ----------
llm = OpenAI(temperature=0, model="gpt-4o-mini", streaming=True)

# ---------- 2. Определяем инструмент ----------
@tool
def get_price(product: str, city: str) -> str:
    """
    Возвращает таблицу цен на указанный продукт в заданном городе.
    """
    # В реальном приложении здесь будет запрос к базе данных или API.
    # Для демонстрации возвращаем статический ответ.
    if product.lower() == "молоко":
        return "| Продукт | Цена (руб.) | Магазин |\n| Молоко  | 89          | Магнит   |"
    elif product.lower() == "хлеб":
        return "| Продукт | Цена (руб.) | Магазин |\n| Хлеб    | 25          | Пятёрочка |"
    else:
        return f"Информация о {product} в городе {city} недоступна."

# ---------- 3. Создаём агента ----------
tools = [get_price]
agent = initialize_agent(
    tools,
    llm,
    agent=AgentType.OPENAI_FUNCTIONS,
    verbose=False,          # отключаем стандартный вывод
)

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

def format_chunk_message(chunk):
    """
    Обрабатываем чанк типа 'messages'.
    Выводим токен без перевода строки, добавляем разделитель при смене шага.
    """
    message, meta = chunk
    global step

    if meta.get("langgraph_step") != step:
        step = meta["langgraph_step"]
        rprint("\n --- --- --- \n")

    content = format_message(message)
    # Выводим без перевода строки, чтобы токены «текли»
    print(content, end="", flush=True)

# ---------- 5. Запускаем потоковый вывод ----------
if __name__ == "__main__":
    # Пример запроса
    user_prompt = "Сколько стоит молоко в Казани? А как насчёт хлеба?"

    stream = agent.stream(
        {"messages": [{"role": "human", "content": user_prompt}]},
        stream_mode=["messages", "updates"],
    )

    step = 1

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

    # Завершающий перевод строки, чтобы консоль оставалась аккуратной
    print()