from langgraph.graph import StateGraph
from typing import TypedDict, List, Any
from rich.console import Console

# Define state schema
class AgentState(TypedDict):
    conversation_history: List[dict]
    last_tool_call_confirmed: bool | None

console = Console()

# Memory node: just returns the existing history
async def memory_node(state: AgentState) -> dict:
    return {"memory": state.get("conversation_history", [])}

# Confirm tool call node
async def confirm_tool_node(state: AgentState, tool_call: Any) -> dict:
    try:
        console.print(f"[bold cyan]Tool call:[/bold cyan] {tool_call}")
        confirmation = input("Confirm execution? (y/n): ")
        confirmed = confirmation.lower().startswith("y")
        state["last_tool_call_confirmed"] = confirmed
    except Exception as e:
        console.print(f"[red]Error during confirmation:[/red] {e}")
        state["last_tool_call_confirmed"] = False
    return {"confirmed_tool_call": state["last_tool_call_confirmed"]}

# Agent node: placeholder for actual LangChain agent logic
async def agent_node(state: AgentState, memory: Any, confirmed_tool_call: Any) -> dict:
    # In a real implementation we would invoke the agent here.
    # For demonstration, echo back the conversation history and confirmation flag.
    response = {
        "memory": memory,
        "confirmed": confirmed_tool_call,
        "message": "Agent executed with memory and confirmation."
    }
    return {"response": response}

# Build graph following execution flow: MemoryNode -> AgentNode -> ConfirmToolNode -> AgentNode
builder = StateGraph(AgentState)
builder.add_node("MemoryNode", memory_node)
builder.add_node("AgentNode", agent_node)
builder.add_node("ConfirmToolNode", confirm_tool_node)

# Define transitions according to plan execution flow
builder.set_entry_point("MemoryNode")
builder.add_edge("MemoryNode", "AgentNode")
builder.add_edge("AgentNode", "ConfirmToolNode")
builder.add_edge("ConfirmToolNode", "AgentNode")
builder.set_finish_node("AgentNode")

# Compile graph
graph = builder.compile()

if __name__ == "__main__":
    # Example initial state
    init_state: AgentState = {"conversation_history": [], "last_tool_call_confirmed": None}
    result = graph.invoke(init_state)
    console.print(result)