"""Stream-режим AI-агента: потоковый вывод токенов и событий updates."""
from __future__ import annotations

import os

from dotenv import load_dotenv
from langchain.agents import create_agent
from langchain.tools import tool
from langchain_openai import ChatOpenAI
from pydantic import SecretStr

load_dotenv()

llm = ChatOpenAI(
    model=os.getenv("OPENAI_MODEL", "openai/gpt-oss-20b:free"),
    base_url=os.getenv("OPENAI_BASE_URL", "https://openrouter.ai/api/v1"),
    api_key=SecretStr(os.getenv("OPENAI_API_KEY", "fake")),
    temperature=0.7,
)

step = 1


@tool
def get_price(product: str, city: str) -> str:
    """Узнать примерную цену продукта в указанном городе. Возвращает строку таблицы."""
    price_agent = create_agent(
        model=llm,
        system_prompt=(
            "Ты эксперт по розничным ценам в России. "
            "Ответ — одна строка таблицы: | Продукт | Цена (руб.) | Магазин |"
        ),
    )
    result = price_agent.invoke(
        {
            "messages": [
                {
                    "role": "human",
                    "content": (
                        f"Какая примерная цена на «{product}» в городе {city}? "
                        "Верни строку | Продукт | Цена (руб.) | Магазин |"
                    ),
                }
            ]
        }
    )
    return result["messages"][-1].content


shopping_agent = create_agent(
    model=llm,
    tools=[get_price],
    system_prompt="Ты помощник по планированию покупок",
)


def format_message(message) -> str:
    """Текст сообщения или вызов инструмента."""
    if message.content:
        return str(message.content)
    tool_calls = getattr(message, "tool_calls", None) or []
    if tool_calls:
        tc = tool_calls[0]
        name = tc.get("name") if isinstance(tc, dict) else getattr(tc, "name", "?")
        args = tc.get("args") if isinstance(tc, dict) else getattr(tc, "args", {})
        return f"{name}({args})"
    return str(message)


def format_chunk_message(chunk_data: tuple) -> None:
    """Потоковый вывод токенов с разделителем при смене шага."""
    global step
    message, meta = chunk_data
    current_step = meta.get("langgraph_step", step)

    if current_step != step:
        step = current_step
        print("\n --- --- --- \n")

    if message.content:
        print(message.content, end="", flush=True)


def main() -> None:
    global step
    step = 1

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

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

    for chunk in stream:
        chunk_type, chunk_data = chunk

        if chunk_type == "messages":
            format_chunk_message(chunk_data)

        if chunk_type == "updates":
            model_update = chunk_data.get("model")
            if model_update:
                last_message = model_update["messages"][-1]
                formatted = format_message(last_message)
                if formatted:
                    print(formatted)

    print()


if __name__ == "__main__":
    main()