"""
Stream-режим AI-агента

Агент с потоковым выводом: замена .invoke() на .stream()
с stream_mode=['messages', 'updates'].
Текст появляется в консоли по мере генерации, с разделителями между шагами.
"""

import os
from rich.console import Console
from langchain_core.tools import tool
from langchain_openai import ChatOpenAI
from langgraph.prebuilt import create_react_agent

console = Console()


@tool
def get_weather(city: str) -> str:
    """Получить текущую погоду в указанном городе."""
    weather_data = {
        "Казань": "🌤 Казань: +18°C, переменная облачность",
        "Москва": "☀️ Москва: +22°C, ясно",
        "Санкт-Петербург": "🌧 Санкт-Петербург: +14°C, дождь",
    }
    return weather_data.get(city, f"🌡 {city}: данные о погоде недоступны")


@tool
def search_info(query: str) -> str:
    """Поиск информации по запросу."""
    return f"Результат поиска по запросу: '{query}' — найдено 5 релевантных источников."


def format_message(message):
    """Форматирование сообщения для вывода."""
    if hasattr(message, "content") and message.content:
        return message.content
    return str(message)


def run_stream_agent():
    """Запуск агента в stream-режиме."""

    llm = ChatOpenAI(
        model=os.environ.get("OPENAI_MODEL", "gpt-4o-mini"),
        openai_api_base=os.environ.get("OPENAI_API_BASE", "https://api.openai.com/v1"),
        openai_api_key=os.environ.get("OPENAI_API_KEY", ""),
        temperature=0,
    )

    agent = create_react_agent(
        model=llm,
        tools=[get_weather, search_info],
        prompt="Ты полезный ассистент. Отвечай на вопросы пользователя.",
    )

    user_input = "Какая погода в Казани и Москве? Также найди информацию о LangGraph."

    console.print("=" * 60)
    console.print(f"  Запрос: {user_input}")
    console.print("=" * 60)
    console.print()

    prev_step = None

    # Заменили .invoke() на .stream() с stream_mode=['messages', 'updates']
    for chunk in agent.stream(
        {"messages": [{"role": "user", "content": user_input}]},
        stream_mode=["messages", "updates"],
    ):
        chunk_type, chunk_data = chunk

        # Обработка 'messages' — токены текста
        if chunk_type == "messages":
            message, meta = chunk_data
            if hasattr(message, "content") and message.content:
                # При смене langgraph_step — печатаем разделитель
                current_step = meta.get("langgraph_step")
                if prev_step is not None and current_step != prev_step:
                    console.print()
                    console.print("--- --- ---")
                    console.print()
                prev_step = current_step

                # Выводим контент сразу без перевода строки
                console.print(message.content, end="", style="white")

        # Обработка 'updates' — события вызова инструментов
        elif chunk_type == "updates":
            if "model" in chunk_data:
                model_output = chunk_data["model"]
                if "messages" in model_output and model_output["messages"]:
                    last_msg = model_output["messages"][-1]
                    formatted = format_message(last_msg)
                    if formatted:
                        console.print(f"\n[bold cyan]Шаг модели:[/bold cyan] {formatted}")

    console.print()
    console.print()
    console.print("=" * 60)
    console.print("  Ответ получен (stream-режим)")
    console.print("=" * 60)


if __name__ == "__main__":
    run_stream_agent()