import os
from langchain_openai import ChatOpenAI
from langgraph.checkpoint.memory import MemorySaver
from langgraph.prebuilt import create_react_agent as create_agent
from rich.console import Console
from rich.markdown import Markdown
from tools import get_price

# LLM configuration using BroJS
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.5,
)

# Memory checkpoint for interrupt handling
memory = MemorySaver()

# Create the agent with interrupt before tools
agent = create_agent(
    model=llm,
    tools=[get_price],
    checkpointer=memory,
    interrupt_before=["tools"],
)

console = Console()


def ask_and_run(user_input: str, config: dict):
    """Synchronously run the agent with streaming and handle interrupts."""
    # Stream messages and updates
    for chunk in agent.stream(
        {"messages": [{"role": "user", "content": user_input}],
         "configurable": config},
        stream_mode=["messages", "updates"],
    ):
        if isinstance(chunk, dict) and "__interrupt__" in chunk:
            # Interrupt: ask for confirmation
            console.print("[bold red]Agent requested tool execution. Confirm? (y/n): ", end="")
            choice = input().strip().lower()
            if choice != "y":
                # Reject by sending a new message to the agent
                config.update({"configurable": {"thread_id": config.get("thread_id", "default")}})
                continue
        if isinstance(chunk, dict) and "messages" in chunk:
            for msg in chunk["messages"]:
                console.print(Markdown(msg["content"]))

if __name__ == "__main__":
    thread_id = os.getenv("THREAD_ID", "session-1")
    config = {"thread_id": thread_id}
    console.print("[bold green]Welcome to the price agent. Type 'exit' to quit.")
    while True:
        user_input = input("Вы: ")
        if user_input.lower() in ("exit", "quit"):
            break
        ask_and_run(user_input, config)
    console.print("[bold blue]Goodbye!", style="bold")