"""
graph.py - MCP Agent Graph for AWS and Microsoft Documentation

This module defines the LangGraph agent for the MCP (Multi-Cloud Platform) assistant,
enabling intelligent routing and answering of queries related to AWS and Microsoft technologies.

Key Features:
- Integrates with AWS and Microsoft MCP servers for documentation and diagram generation.
- Dynamically loads available tools from both AWS and Microsoft MCP endpoints.
- Provides a unified agent that can answer questions about AWS services, Microsoft technologies, and generate AWS diagrams.
- Uses OpenAI GPT-4o-mini as the LLM backend for reasoning and tool selection.
- Returns expert-level answers or defers when information is insufficient.

Main Functions:
- load_tools(): Asynchronously loads tool definitions from AWS and Microsoft MCP servers.
- build_graph(): Constructs and returns the LangGraph agent configured with the appropriate tools and system prompt.

Intended Usage:
Import and call `build_graph()` to obtain a ready-to-use agent for enterprise architecture, cloud, and documentation queries.

References:
- [AWS MCP Servers](https://github.com/awslabs/mcp/tree/main/src)
- [MicrosoftDocs MCP Server](https://github.com/microsoftdocs/mcp)
"""
import asyncio
import nest_asyncio

from langchain_mcp_adapters.client import MultiServerMCPClient
from langgraph.prebuilt import create_react_agent

from ea4all.src.ea4all_mcp.configuration import AgentConfiguration
from ea4all.src.ea4all_mcp.prompts import MCP_AWS_MS_SYSTEM_PROMPT

from ea4all.src.shared.utils import (
    get_llm_client, 
)

async def load_tools():
    # Set up the MCP client for AWS Documentation
    mcp_client = MultiServerMCPClient(
        {
            "awslabs.aws-documentation-mcp-server": {
                "transport": "stdio",
                "command": "uvx",
                "args": ["awslabs.aws-documentation-mcp-server@latest"],
                "env": {
                    "FASTMCP_LOG_LEVEL": "ERROR",
                    "AWS_DOCUMENTATION_PARTITION": "aws"
                },
                "disabled": False,
                "autoApprove": [],
            },
            "awslabs.aws-diagram-mcp-server": {
                "transport": "stdio",
                "command": "uvx",
                "args": ["awslabs.aws-diagram-mcp-server"],
                "env": {
                    "FASTMCP_LOG_LEVEL": "ERROR"
                },
                "autoApprove": [],
                "disabled": False
            },
            "microsoft_mcp_tools": {
                "url": "https://learn.microsoft.com/api/mcp",
                "transport": "streamable_http",
            }
        }
    )

    # Fetch the tools from the AWS MCP server
    mcp_tools = await mcp_client.get_tools()

    return mcp_tools

async def build_graph():
    """Build the LangGraph graph for the MCP agent."""
    # Define the graph structure
    system_prompt = """
    """

    mcp_tools = await load_tools()
    
    llm = get_llm_client(model=AgentConfiguration.mcp_docs_model)

    # Create the LangGraph agent for AWS documentation
    ea4all_docs = create_react_agent(
        model=llm,
        prompt=MCP_AWS_MS_SYSTEM_PROMPT,
        tools=mcp_tools,
        debug=True,
        #state_schema=OverallState,
    )

    ea4all_docs.name = "ea4all_mcp_aws_microsoft_docs_agent"

    return ea4all_docs

nest_asyncio.apply()
mcp_docs_graph = asyncio.run(build_graph())