import os
from langchain_openai import ChatOpenAI
from langgraph.checkpoint.memory import MemorySaver
from langgraph.prebuilt import create_react_agent
from rich.console import Console

# ---------------------------------------------------------------
# LLM
# ---------------------------------------------------------------
llm = ChatOpenAI(
    model="openai/gpt-oss-20b:free",
    base_url="https://platform.brojs.ru/jrnl-bh/api/inference/v1",
    api_key=os.getenv("JOURNAL_MCP_PAT"),
    temperature=0.0,
)

# ---------------------------------------------------------------
# Helper create_agent — обёртка, как описано в инструкции
# ---------------------------------------------------------------
def create_agent(model, tools, system_prompt="", checkpointer=None, interrupt_before=None):
    """Helper-функция для создания агента с памятью и interrupt."""
    kwargs = dict(model=model, tools=tools)
    if system_prompt:
        kwargs["prompt"] = system_prompt
    if checkpointer is not None:
        kwargs["checkpointer"] = checkpointer
    if interrupt_before is not None:
        kwargs["interrupt_before"] = interrupt_before
    return create_react_agent(**kwargs)

# ---------------------------------------------------------------
# Инструменты
# ---------------------------------------------------------------
from langchain.tools import tool

@tool
def get_price(city: str, date: str) -> str:
    """Возвращает условную цену/погоду для города на указанную дату."""
    return f"Данные для {city} на {date}: 720 руб/кг, 18C, облачно"

# ---------------------------------------------------------------
# Агент с памятью и паузой перед инструментом
# ---------------------------------------------------------------
memory = MemorySaver()

agent = create_agent(
    model=llm,
    tools=[get_price],
    system_prompt="Ты полезный ассистент. Помогай пользователю находить цены и погоду.",
    checkpointer=memory,
    interrupt_before=["tools"],
)

console = Console()

# ---------------------------------------------------------------
# ask_and_run — точная структура из задания
# ---------------------------------------------------------------
def ask_and_run(user_input, config):
    for chunk in agent.stream(user_input, config=config, stream_mode=["messages", "updates"]):
        state = agent.get_state(config)
        chunk_type, chunk_data = chunk

        if chunk_type == "messages":
            msg = chunk_data[0] if isinstance(chunk_data, (list, tuple)) else chunk_data
            content = getattr(msg, "content", "")
            if content and not getattr(msg, "tool_calls", None):
                console.print(content, end="")

        if chunk_type == "updates":
            console.print("\n --- --- --- \n")
            for node_data in chunk_data.values():
                for msg in node_data.get("messages", []):
                    for tc in getattr(msg, "tool_calls", []):
                        console.print(f"[bold]{tc['name']}({tc['args']})[/bold]")

        if "__interrupt__" in chunk_data and state.next == ("tools",):
            tool_call = state.values["messages"][-1].tool_calls[0]
            console.print(f"Агент хочет вызвать утилиту [bold]{tool_call['name']}({tool_call['args']})[/bold]")
            answer = input("Разрешить? (Y/n): ")
            if answer.lower().strip() == "y":
                ask_and_run(None, config)
            else:
                console.print("Отменено")
                break

# ---------------------------------------------------------------
# Чат-цикл
# ---------------------------------------------------------------
if __name__ == "__main__":
    config = {"configurable": {"thread_id": "разговор-1"}}
    console.print("[bold green]Агент запущен. Введите 'exit' для выхода.[/bold green]")

    while True:
        user_input = input("\nВы: ")
        if user_input == "exit":
            break

        ask_and_run(
            {"messages": [{"role": "human", "content": user_input}]},
            config,
        )