from typing import TypedDict, Annotated, Sequence
from langchain_core.messages import BaseMessage, AIMessage
from langgraph.graph import StateGraph, END
from src.agent import create_agent_executor

class AgentState(TypedDict):
    messages: Annotated[Sequence[BaseMessage], lambda x, y: x + y]

agent_executor = create_agent_executor()

async def run_agent_executor(state: AgentState):
    """
    Runs the agent executor and returns the FINAL output message,
    correctly wrapped in an AIMessage object.
    """
    print("--- [Graph] Running Agent Executor ---")
    inputs = {
        "input": state['messages'][-1].content,
        "chat_history": state['messages'][:-1],
    }
    
    agent_outcome = await agent_executor.ainvoke(inputs)
    
    final_response_string = agent_outcome["output"]
    return {"messages": [AIMessage(content=final_response_string)]}

def create_graph(checkpointer):
    """
    Creates a simple LangGraph with a single node that runs our agent.
    """
    workflow = StateGraph(AgentState)

    workflow.add_node("agent", run_agent_executor)
    workflow.set_entry_point("agent")
    workflow.add_edge("agent", END)

    app = workflow.compile(checkpointer=checkpointer)
    return app