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, мы просто продолжаем выполнение (возобновление после паузы)
    input_payload = {"messages": [{"role": "human", "content": user_input}]} if user_input else None
    
    # stream_mode=['messages', 'updates'] позволяет обрабатывать и токены, и изменения состояния
    stream = agent.stream(input_payload, config=config, stream_mode=['messages', 'updates'])
    
    try:
        for chunk in stream:
            chunk_type, chunk_data = chunk
            state = agent.get_state(config)

            if chunk_type == 'messages':
                # Потоковый вывод текста (токены)
                # В новых версиях LangGraph/LangChain чанк сообщения может быть объектом BaseMessage
                content = ""
                if hasattr(chunk_data, 'content'):
                    content = chunk_data.content
                elif isinstance(chunk_data, str):
                    content = chunk_data
                
                if content:
                    console.print(content, end="", flush=True)

            elif chunk_type == 'updates':
                # Здесь можно обрабатывать обновления узлов, если нужно
                pass

            # Проверка на прерывание (интеррупт перед инструментом)
            # В LangGraph при interrupt_before ключ '__interrupt__' появляется в данных чанка
            if '__interrupt__' in chunk_data and state.next == ('tools',):
                console.print("\n" + "---" * 10)
                
                # Извлекаем информацию о том, какой инструмент вызывается
                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':
                    # Рекурсивный вызов с None для продолжения выполнения
                    ask_and_run(None, config)
                else:
                    console.print("Отменено")
                    return # Выход из текущей итерации функции

        # Печать переноса строки после завершения стрима (если не было прерывания)
        if user_input is not None:
            console.print() 

    except Exception as e:
        # Обработка ошибок или остановки потока
        pass

def main():
    console.print("[bold blue]Чат запущен. Напишите 'exit' для выхода.[/bold blue]")
    
    while True:
        try:
            user_text = input("\nВы: ")
            if user_text.lower().strip() == 'exit':
                console.print("До свидания!")
                break

            # Запуск агента с новым сообщением
            ask_and_run(
                {"messages": [{"role": "human", "content": user_text}]},
                config,
            )
        except EOFError:
            break
        except Exception as e:
            console.print(f"Ошибка: {e}")

if __name__ == "__main__":
    main()