import os
from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage, SystemMessage
from langchain.tools import tool
from langgraph.checkpoint.memory import MemorySaver
from langgraph.types import Command
from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages
from rich.console import Console

# LLM setup
llm = ChatOpenAI(
    model="openai/gpt-oss-20b:free",
    base_url="https://openrouter.ai/api/v1",
    api_key=os.getenv("OPENAI_API_KEY"),
    temperature=0.0,
)

console = Console()

# Dummy tool – can be replaced by real implementation
@tool
def get_price(query: str) -> str:
    """Return a fake price for demonstration."""
    return f"Цена для {query} – 100₽"

# State definition
class ChatState(dict):
    messages: list = add_messages
    location: str = "start"
    inventory: list = []

# Agent graph
builder = StateGraph(ChatState)

# Node that runs the LLM and returns tool calls
async def llm_node(state: ChatState):
    system = SystemMessage(content="You are a helpful assistant. Use the get_price tool if needed.")
    response = await llm.ainvoke([system] + state["messages"], config={"configurable": {"thread_id": "chat-1"}})
    return {"messages": [response]}

# Node that handles tool calls (only get_price)
async def tool_node(state: ChatState):
    last_msg = state["messages"][-1]
    if hasattr(last_msg, "tool_calls") and last_msg.tool_calls:
        tool_call = last_msg.tool_calls[0]
        name = tool_call["name"]
        args = tool_call["args"]
        if name == "get_price":
            result = get_price(args)
            # Append tool result to messages
            state["messages"].append(HumanMessage(content=result))
    return state

builder.add_node("llm", llm_node)
builder.add_node("tool", tool_node)
builder.add_edge(START, "llm")
builder.add_edge("llm", "tool")
builder.add_edge("tool", END)

chat = builder.compile(checkpointer=MemorySaver(), interrupt_before=["tool"])

# Main loop
while True:
    user_input = console.input("\nВы: ")
    if user_input.lower() in {"exit", "quit"}:
        break
    # Add user message
    chat.add_state({"messages": [HumanMessage(content=user_input)]})
    # Run graph with interrupt handling
    state = chat.get_state({"configurable": {"thread_id": "chat-1"}})
    while True:
        result = await chat.ainvoke(state, {"configurable": {"thread_id": "chat-1"}})
        # If interrupted before tool
        if "__interrupt__" in result and result["__interrupt__"][0]["name"] == "tool":
            console.print("\nАгент хочет вызвать инструмент: get_price")
            ans = console.input("Разрешить? (Y/n): ")
            if ans.lower() in {"", "y", "yes"}:
                # resume
                state = chat.get_state({"configurable": {"thread_id": "chat-1"}})
                continue
            else:
                console.print("Отменено")
                break
        # Print assistant response
        for msg in result["messages"]:
            if msg.role == "assistant":
                console.print(msg.content)
        break

console.print("\nДо свидания!")