"""
Практическое задание: Stream-режим AI-агента
Заменить единовременный вызов .invoke() на потоковый вывод через .stream(),
чтобы ответ появлялся в консоли по мере генерации токен за токеном.
"""

import os
from typing import Literal
from langchain_openai import ChatOpenAI
from langchain.agents import create_agent, AgentExecutor
from langchain_core.tools import tool
from langgraph.checkpoint.memory import MemorySaver

# Настройка API ключа
os.environ["OPENAI_API_KEY"] = os.getenv("OPENAI_API_KEY", "your-api-key-here")

# Пример инструмента
@tool
def get_weather(city: str) -> str:
    """Получить погоду в городе"""
    weather_data = {
        "Казань": "В Казани сегодня +15°C, облачно",
        "Москва": "В Москве сегодня +12°C, дождь",
        "Санкт-Петербург": "В Санкт-Петербурге сегодня +10°C, солнечно",
    }
    return weather_data.get(city, f"Погода в {city} неизвестна")


@tool
def get_time(city: str) -> str:
    """Получить текущее время в городе"""
    return f"Сейчас в {city} 14:30 (примерное время)"


# Создание агента с MemorySaver
memory = MemorySaver()
llm = ChatOpenAI(model="gpt-4o", temperature=0)

agent = create_agent(
    model=llm,
    tools=[get_weather, get_time],
    system_prompt='Ты полезный ассистент. Отвечай на вопросы о погоде и времени.',
    checkpointer=memory,
)

# Конфигурация с thread_id
config = {"configurable": {"thread_id": "test-thread-1"}}

def run_stream_agent(question: str):
    """
    Запускает агента в stream-режиме.
    Обрабатывает чанки типов 'messages' и 'updates'.
    """
    print(f"\nВы: {question}")
    print("-" * 40)
    
    # Потоковый вызов агента
    for chunk_type, chunk_data in agent.stream(
        {"messages": [("user", question)]},
        config,
        stream_mode=['messages', 'updates']
    ):
        if chunk_type == 'messages':
            # Потоковый вывод токенов без переноса строки
            if hasattr(chunk_data, 'content'):
                print(chunk_data.content, end="", flush=True)
        elif chunk_type == 'updates':
            # Вывод событий об изменении состояния
            if chunk_data.get('langgraph_step'):
                print("\n" + "-" * 40)
                step = chunk_data['langgraph_step']
                print(f"[Шаг: {step}]")
            # Вывод информации о вызове инструмента
            if 'tool_calls' in chunk_data:
                for call in chunk_data['tool_calls']:
                    print(f"\nВызов инструмента: {call.get('name')}")
                    print(f"Аргументы: {call.get('args')}")
    
    print("\n" + "=" * 40)


def main():
    """Главная функция для демонстрации stream-режима."""
    print("=" * 40)
    print("Stream-режим AI-агента")
    print("=" * 40)
    
    questions = [
        "Какая погода в Казани?",
        "Сколько времени в Москве?",
        "Какая погода будет в Санкт-Петербурге?"
    ]
    
    for q in questions:
        run_stream_agent(q)


if __name__ == "__main__":
    main()