from langchain_openai import ChatOpenAI
from pydantic import SecretStr
from langchain.agents import create_agent
from langchain.tools import tool
from langgraph.checkpoint.memory import MemorySaver
from rich.console import Console

# Инициализация LLM с использованием плейсхолдеров
llm = ChatOpenAI(
    model="google/gemma-4-26b-a4b",
    base_url="http://192.168.0.120:1234/v1",
    api_key=SecretStr("lm-studio"),
    temperature=0.7,
)

console = Console()

# Определение инструмента
@tool
def get_price(city: str, date: str):
    """Возвращает прогноз погоды (цены/состояние) для указанного города и даты."""
    # Имитация логики
    return f"В городе {city} на дату {date} ожидается солнечная погода, +20°C."

tools = [get_price]

# Настройка памяти и агента с механизмом прерывания (interrupt)
memory = MemorySaver()
agent = create_agent(
    model=llm,
    tools=tools,
    system_prompt="Ты полезный помощмущник. Если пользователь спрашивает о погоде, используй инструмент get_price.",
    checkpointer=memory,
    interrupt_before=['tools'],
)

# Конфигурация потока (thread_id обеспечивает память разговора)
config = {"configurable": {"thread_id": "chat-session-123"}}

def ask_and_run(user_input, config):
    """Основная функция обработки сообщений и управления циклом подтверждения."""
    
    # Если user_input is None, мы просто продолжаем выполнение (возобновление после паузы)
    # В LangGraph для возобновления через stream передается None или пустой список сообщений
    stream_input = user_input if user_input is not None else []
    
    # Используем stream_mode=['messages', 'updates'] согласно заданию
    try:
        for chunk in agent.stream(stream_input, config=config, stream_mode=['messages', 'updates']):
            state = agent.get_state(config)
            chunk_type, chunk_data = chunk

            if chunk_type == 'messages':
                # Потоковый вывод токенов (сообщений)
                message, metadata = chunk_data
                if hasattr(message, "content") and message.content:
                    print(message.append if hasattr(message, 'append') else message.content, end="", flush=True)

            elif chunk_type == 'updates':
                # Вывод информации об обновлениях (например, вызовы инструментов)
                pass

            # Проверка на прерывание перед вызовом инструмента
            if '__interrupt__' in chunk_data and state.next == ('tools',):
                console.print("\n" + "---" * 15)
                
                # Извлекаем информацию о том, какой инструмент хочет вызвать агент
                last_message = state.values['messages'][-1]
                if hasattr(last_message, 'tool_calls') and last_message.tool_calls:
                    tool_call = last_message.tool_calls[0]
                    console.print(f"{tool_call['name']}({tool_call['args']})")
                    console.print(f"Агент хочет вызвать утилиту {tool_call['name']}({tool_call['args']})")
                
                answer = input("Разрешить? (Y/n): ")
                
                if answer.lower().strip() == 'y':
                    # Рекурсивный вызов для продолжения без нового сообщения пользователя
                    ask_and_run(None, config)
                else:
                    console.print("Отменено")
                    return # Выходим из текущей итерации стрима

        if user_input: # Печать переноса строки только если был новый ввод
            print()
            
    except Exception as e:
        # Если произошла ошибка в процессе стрима (например, при рекурсии), пробрасываем её выше
        raise e

if __name__ == "__main__":
    console.print("[bold blue]Чат запущен. Напишите 'exit' для выхода.[/bold blue]")
    
    while True:
        try:
            user_text = input("\nВы: ")
            if user_text.lower().strip() == 'exit':
                break
            
            # Формируем входные данные для агента
            input_payload = {"messages": [{"role": "human", "content": user_text}]}
            
            ask_and_run(input_payload, config)
            
        except KeyboardInterrupt:
            break
        except Exception as e:
            console.print(f"[bold red]Ошибка: {e}[/bold red]")