# agent.py

"""
Простой LangChain‑агент с поддержкой потокового вывода через `agent.stream()`.
Внутри используется один инструмент – get_price, который вызывает под‑агента,
обрабатывающего запросы к «базе цен».
"""

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

# --------------------------------------------------------------------------- #
# 1. Определяем LLM (Ollama) – это единственный доступный LLM в задании.
# --------------------------------------------------------------------------- #
try:
    from langchain_ollama import ChatOllama
except ImportError:  # если Ollama недоступен, используем простой Mock‑LLM
    class DummyLLM:
        def __call__(self, *args, **kwargs):
            return "Mock response"

    llm = DummyLLM()
else:
    llm = ChatOllama(model="llama3.1")  # пример модели

# --------------------------------------------------------------------------- #
# 2. Создаём инструмент get_price.
# --------------------------------------------------------------------------- #

@tool
def get_price(product: str, city: str) -> str:
    """
    Возвращает таблицу цен на указанный продукт в заданном городе.
    Для демонстрации используется под‑агент, который формирует ответ.
    """
    # Под‑агент – простая функция, которая возвращает строку
    sub_agent = create_agent(
        llm=llm,
        tools=[],
        verbose=False,
        agent_type="openai-tools",
    )

    # Запускаем под‑агента в режиме invoke (т.к. он короткий)
    result = sub_agent.invoke({
        "messages": [
            {"role": "human", "content": f"Сколько стоит {product} в городе {city}?"}
        ]
    })
    return result["output"]

# --------------------------------------------------------------------------- #
# 3. Создаём главный агент с инструментом get_price.
# --------------------------------------------------------------------------- #

agent = create_agent(
    llm=llm,
    tools=[get_price],
    verbose=False,
    agent_type="openai-tools",
)

# --------------------------------------------------------------------------- #
# 4. Функции для форматирования и вывода чанков.
# --------------------------------------------------------------------------- #

step = 1

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

def format_message(message):
    """
    Форматирует сообщение для вывода в режиме invoke.
    """
    if message.content:
        return message.content
    # Если нет content – это вызов инструмента
    tool_call = message.tool_calls[0]
    return f"{tool_call['name']}({tool_call['args']})"

# --------------------------------------------------------------------------- #
# 5. Запускаем потоковый вывод.
# --------------------------------------------------------------------------- #

if __name__ == "__main__":
    # Пример запроса к агенту
    stream = agent.stream(
        {
            "messages": [
                {"role": "human", "content": "Покажи мне цены на молоко и хлеб в Казани."}
            ]
        },
        stream_mode=["messages", "updates"],
    )

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