import sys
from langgraph.checkpoint.memory import MemorySaver
from langgraph.graph import StateGraph
from langchain_openai import ChatOpenAI
from langchain.tools import tool
from rich.console import Console

console = Console()

# Simple tool example
@tool
def echo(text: str) -> str:
    """Return the same text."""
    return text

# Define graph state
class State(dict):
    pass

# Agent node
async def agent_node(state: State, config=None):
    # Use LangChain LLM with tool calling
    llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)
    from langchain.agents import create_openai_functions_agent
    agent = create_openai_functions_agent(llm, [echo])
    messages = state.get("messages", [])
    # Run agent
    result = await agent.ainvoke({"messages": messages})
    return {"messages": result["messages"]}

# Build graph
builder = StateGraph(State)
builder.add_node("agent", agent_node)
builder.set_entry_point("agent")

graph = builder.compile(checkpointer=MemorySaver(), interrupt_before=["tools"])  # memory and pause before tools

config = {"configurable": {"thread_id": "chat-1"}}

async def run_chat():
    import asyncio
    while True:
        user_input = input("Вы: ")
        if user_input.lower() in ("exit", "quit"):
            break
        state = {"messages": [{"role": "user", "content": user_input}]}
        # Stream output and handle pauses
        async for chunk in graph.astream(state, config=config):
            if isinstance(chunk, dict) and chunk.get("__interrupt__"):
                # Pause before tool call
                console.print("\n[bold yellow]Agent wants to call a tool. Confirm? (y/n)\b")
                ans = input()
                if ans.lower() != "y":
                    console.print("[red]Cancelled by user.[/]")
                    break
            else:
                # Print token stream
                console.print(chunk.get("messages", [])[0].get("content", ""), end="")
        console.print()

if __name__ == "__main__":
    import asyncio
    asyncio.run(run_chat())