import asyncio
from pathlib import Path

from llama_index.core.agent.workflow import FunctionAgent, ToolCallResult
from llama_index.core.base.llms.types import ChatMessage, MessageRole, ToolCallBlock
from llama_index.core.llms.mock import MockFunctionCallingLLM
from llama_index.tools.mcp import BasicMCPClient, McpToolSpec


SERVER_PATH = Path(__file__).with_name("release_tools_server.py")
QUESTION = "Check deployment ticket SR-431."


def mcp_result_text(tool_output) -> str:
    raw_output = getattr(tool_output, "raw_output", None)
    structured = getattr(raw_output, "structuredContent", None)
    if isinstance(structured, dict) and structured.get("result"):
        return str(structured["result"])

    content = getattr(raw_output, "content", None)
    if content:
        text = getattr(content[0], "text", None)
        if text:
            return str(text)

    return str(tool_output)


def response_generator(messages, **kwargs):
    tool_messages = [
        str(message.content or "")
        for message in messages
        if message.role == MessageRole.TOOL
    ]
    if tool_messages:
        return ChatMessage(
            role=MessageRole.ASSISTANT,
            content=(
                "The deployment ticket SR-431 is approved for production release."
            ),
        )

    return ChatMessage(
        role=MessageRole.ASSISTANT,
        blocks=[
            ToolCallBlock(
                tool_call_id="call_lookup_ticket_1",
                tool_name="lookup_ticket",
                tool_kwargs={"ticket_id": "SR-431"},
            )
        ],
    )


async def main() -> None:
    mcp_client = BasicMCPClient("python3", args=[str(SERVER_PATH)])
    tool_spec = McpToolSpec(client=mcp_client, allowed_tools=["lookup_ticket"])
    tools = await tool_spec.to_tool_list_async()

    agent = FunctionAgent(
        tools=tools,
        llm=MockFunctionCallingLLM(
            response_generator=response_generator,
            is_chat_model=True,
        ),
        system_prompt="Use MCP tools for deployment ticket questions.",
        streaming=False,
    )

    print("mcp_tools=" + ",".join(tool.metadata.name for tool in tools))
    handler = agent.run(QUESTION)

    async for event in handler.stream_events():
        if isinstance(event, ToolCallResult):
            print(f"tool_called={event.tool_name}")
            print(f"tool_input={event.tool_kwargs['ticket_id']}")
            print(f"tool_output={mcp_result_text(event.tool_output)}")

    response = await handler
    print(f"agent_answer={response}")


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