# MCPClient is a Pit
# MCPClient is a service that gets tools information from the Anthropic Model Content Protocol service
# and creates practices for the agent

from prompits.Pit import Pit
from mcp import ClientSession, StdioServerParameters, types
from mcp.client.stdio import stdio_client
import asyncio

from prompits.Practice import Practice


class MCPClient(Pit):
    def __init__(self, name, description, mcp_server_params: StdioServerParameters):
        super().__init__(name, description)
        self.mcp_server_params = mcp_server_params
        self.tools= asyncio.run(self._GetTools())
        self.AddPractice(Practice("GetTools", self._GetTools, 
                                  description="Get tools from the MCP server", 
                                  is_async=True))
        for tool in self.tools:
            #print(tool)
            self.AddPractice(Practice(tool.name, 
                                      function= self._CallTool, input_schema=tool.inputSchema, 
                                      description=tool.description, parameters={"tool_name":tool.name},
                                      is_async=True))
        #print(self.practices)
        #self.mcp_client = ClientSession(mcp_server_params)

    # Optional: create a sampling callback
    async def handle_sampling_message(
        message: types.CreateMessageRequestParams,
    ) -> types.CreateMessageResult:
        result= types.CreateMessageResult(
            role="assistant",
            content=types.TextContent(
                type="text",
                text="Hello, world! from model",
            ),
            model="gpt-3.5-turbo",
            stopReason="endTurn",
        )
        print(f"handle_sampling_message: {result}")
        return result
    
    async def _CallTool(self, tool_name, arguments):
        async with stdio_client(self.mcp_server_params) as (read, write):
            async with ClientSession(
                read, write, sampling_callback=self.handle_sampling_message
            ) as session:
                # Initialize the connection
                await session.initialize()

                return await session.call_tool(tool_name, arguments)
    
    async def _GetTools(self):
        async with stdio_client(self.mcp_server_params) as (read, write):
            async with ClientSession(
                read, write, sampling_callback=self.handle_sampling_message
            ) as session:
                # Initialize the connection
                await session.initialize()

                # # List available prompts
                # prompts = await session.list_prompts()

                # # Get a prompt
                # prompt = await session.get_prompt(
                #     "example-prompt", arguments={"arg1": "value"}
                # )

                # List available resources
                # resources = await session.list_resources()

                # List available tools
                tools = await session.list_tools()
                for key, value in tools:
                    if key == "tools":
                        return value
                return []
            
                # # Read a resource
                # content, mime_type = await session.read_resource("file://some/path")

                # # Call a tool
                # result = await session.call_tool("tool-name", arguments={"arg1": "value"})

    def FromJson(self, json_data):
        self.mcp_server_params = StdioServerParameters(**json_data)
        self.mcp_client = ClientSession(self.mcp_server_params)

    def ToJson(self):
        return self.mcp_server_params.to_json()
