from langgraph.graph import StateGraph, START, END
from langchain_core.runnables import RunnableConfig
from langgraph.types import Command

from langchain_core.messages import HumanMessage, AIMessage
from typing import List, Dict, Any
from typing_extensions import TypedDict
from typing import Annotated, Literal
from langgraph.graph.message import add_messages
from langgraph.prebuilt import create_react_agent
from langchain_openai import ChatOpenAI
import asyncio


class State(TypedDict):
    messages: Annotated[list, add_messages]


def init_llm():
    return ChatOpenAI(
        base_url="http://127.0.0.1:8000/v1", # mock local llm server
        model="qwen3-32b", # let's assume this is a valid model name
        api_key="none",
    )

def get_book_list() -> List[Dict[str, Any]]:
    """
    Returns a list of available books.

    Returns:
        List[Dict[str, Any]]: A list of books, each containing the book name, author.
    """
    # mock
    return [
        {"author": "abc", "name": "abc"},
        {"author": "def", "name": "def"},
        {"author": "ghi", "name": "ghi"},
        {"author": "jkl", "name": "jkl"},
        {"author": "mnq", "name": "mnq"}
    ]

async def agent_node(
    state: State, config: RunnableConfig
) -> Command[Literal["__end__"]]:
    # mcp_servers = config["configurable"]["mcp_servers"]
    # client = MultiServerMCPClient(mcp_servers)
    # tools = await client.get_tools()

    agent = create_react_agent(
        name="info-retrieval-agent",
        model=init_llm(),
        tools=[get_book_list],
        prompt="{CURRENT_TIME}\n\n You are a helpful e-book library assistant, you can answer the readers' question, such as `the number of books you have`.",
    )

    message = state.get("messages")[-1].content
    # e.g. how many books do you have?
    agent_input = {"messages": [HumanMessage(message)]}
    # raise error here, invalid tool
    result = await agent.ainvoke(input=agent_input, config={"recursion_limit": 5})

    return Command(
        update={"messages": [AIMessage(result["messages"][-1].content)]}, goto="__end__"
    )


async def main():
    builder = StateGraph(State)
    builder.add_node("agent", agent_node)
    builder.add_edge(START, "agent")
    builder.add_edge("agent", END)
    compiled_graph = builder.compile()
    user_input = {
        "messages": "how many books do you have",
    }

    # it will treat non-argument tool -> invalid tool which is not expected
    async for _, _, event_data in compiled_graph.astream(
        user_input,
        config={
            # "mcp_servers": {
            #     "mock": {
            #         "url": "http://127.0.0.1:3000/sse",
            #         "transport": "sse",
            #         "args": [],
            #         "command": [],
            #         "env": [],
            #     }
            # }
        },
        stream_mode=["messages", "updates"],
        subgraphs=True,
    ):
        pass

    # works fine
    # result = await compiled_graph.ainvoke(user_input)


if __name__ == "__main__":
    asyncio.run(main())