# client.py

"""
Интерактивный CLI для LangChain‑агента с потоковым выводом.
Используем только стек из условия:
    - langchain.agents.create_agent
    - langchain.tools.tool
    - langchain_ollama.ChatOllama
    - rich.print (rprint)
"""

from langchain.agents import create_agent
from langchain.tools import tool
from langchain_ollama import ChatOllama
from rich import print as rprint

# ------------------------------------------------------------------
# 1. Определяем LLM‑обёртку – Ollama
# ------------------------------------------------------------------
llm = ChatOllama(model="llama3")          # модель, которую можно заменить

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

# ------------------------------------------------------------------
# 3. Создаём агента, который использует LLM и инструмент get_price
# ------------------------------------------------------------------
agent = create_agent(
    llm=llm,
    tools=[get_price],
    verbose=False,            # будем выводить вручную
)

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

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

# ------------------------------------------------------------------
# 5. Основной цикл CLI
# ------------------------------------------------------------------
def main():
    """
    Запускаем интерактивный режим: пользователь вводит запрос,
    агент отвечает в режиме stream.
    """
    rprint("[bold cyan]LangChain Stream‑Agent[/bold cyan]")
    rprint("Введите ваш вопрос (Ctrl+C для выхода):")
    while True:
        try:
            user_input = input("\n> ")
        except KeyboardInterrupt:
            rprint("\n[red]Выход из программы.[/red]")
            break

        if not user_input.strip():
            continue

        # Запускаем потоковый вывод
        stream = agent.stream(
            {"messages": [{"role": "human", "content": user_input}]},
            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":
                # При завершении шага выводим итоговое сообщение
                model_update = chunk_data.get("model")
                if model_update:
                    last_msg = model_update["messages"][-1]
                    rprint("\n" + format_message(last_msg))

        # После завершения потока добавляем перевод строки для следующего запроса
        print()

if __name__ == "__main__":
    main()