import asyncio
from contextlib import AsyncExitStack
import os
from typing import Optional
from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client
from mcp.client.streamable_http import streamablehttp_client

from langchain_mcp_adapters.tools import load_mcp_tools
from langgraph.prebuilt import create_react_agent
from langchain_openai import ChatOpenAI
from langchain_core.runnables import RunnableConfig
from langchain_core.prompts import ChatPromptTemplate


class MCPClient:
    def __init__(self):
        # Initialize session and client objects
        self.session: Optional[ClientSession] = None
        self.exit_stack = AsyncExitStack()

    async def connect_to_running_server(self, server_url: str = "http://localhost:8000/mcp"):
        """Connect to an already running MCP server via HTTP

        Args:
            server_url: URL of the running MCP server (default: http://localhost:8000/mcp)
        """
        print(f"Connecting to running MCP server at {server_url}")
        
        # Create HTTP transport to connect to running server
        http_transport = await self.exit_stack.enter_async_context(
            streamablehttp_client(
                server_url,
                headers={
                    "authorization": "73a8ef8aa2fe4fd00be57131f79ed0996b2176cc8ab93d5c01451e8fd4a60f1d"
                }
            )
        )
        read_stream, write_stream, get_session_id = http_transport
        self.session = await self.exit_stack.enter_async_context(
            ClientSession(read_stream, write_stream)
        )
        
        await self.session.initialize()
        print("Connected to running MCP server successfully!")

    async def connect_to_new_server(self):
        """Connect to an MCP server by spawning a new process (old method)

        Args:
            server_script_path: Path to the server script (.py or .js)
        """
        print("Spawning new MCP server process...")
        server_params = StdioServerParameters(
            command="E:/anaconda/install/envs/mcp_server/python.exe",
            # cwd="E:/workspace/bubobot-mcp",
            # args=["-m", "bubobot_mcp.main", "--transport", "stdio"],
            args=["-m", "bubobot_mcp", "--transport", "stdio"],
            env={
                "BEARER_TOKEN": "73a8ef8aa2fe4fd00be57131f79ed0996b2176cc8ab93d5c01451e8fd4a60f1d",
            }
        )

        stdio_transport = await self.exit_stack.enter_async_context(stdio_client(server_params))
        self.stdio, self.write = stdio_transport
        self.session = await self.exit_stack.enter_async_context(ClientSession(self.stdio, self.write))

        await self.session.initialize()
        print("Connected to new MCP server process successfully!")

    async def get_tools(self):
        return await load_mcp_tools(self.session)

    async def call_tool(self, tool_name: str, arguments: dict = None):
        """Call a specific MCP tool directly"""
        print("Args dict: ", arguments)
        if arguments is None:
            arguments = {}
            
        # Don't convert to kwargs format - let the session handle it directly
        return await self.session.call_tool(tool_name, arguments)
    
    async def close(self):
        """Close the client connection"""
        await self.exit_stack.aclose()


async def query_with_running_server(text: str, server_url: str = "http://localhost:8000/mcp"):
    """Query using a running MCP server (recommended method)"""
    client = MCPClient()
    try:
        await client.connect_to_running_server(server_url)
        
        # Patch the session's call_tool method to inject authentication
        original_call_tool = client.session.call_tool
        
        async def authenticated_call_tool(tool_name: str, arguments: dict = None):
            if arguments is None:
                arguments = {}
            
            # Handle the kwargs wrapping that MCP server expects
            if "kwargs" in arguments:
                # Parse kwargs if it's a string
                if isinstance(arguments["kwargs"], str):
                    import json
                    try:
                        if arguments["kwargs"]:  # Non-empty string
                            kwargs_dict = json.loads(arguments["kwargs"])
                        else:  # Empty string
                            kwargs_dict = {}
                    except json.JSONDecodeError:
                        kwargs_dict = {}
                else:
                    kwargs_dict = arguments["kwargs"] if isinstance(arguments["kwargs"], dict) else {}
                
                # Add authentication to the kwargs dict
                kwargs_dict["organizationTeamId"] = 46
                
                # Re-encode as JSON string
                arguments["kwargs"] = json.dumps(kwargs_dict)
            else:
                # Direct arguments - add auth params
                arguments["organizationTeamId"] = 46
            
            
            # Call the original method
            return await original_call_tool(tool_name, arguments)
        
        # Replace the session's call_tool method
        client.session.call_tool = authenticated_call_tool
        
        # Now load the tools - they will use our patched session
        original_tools = await load_mcp_tools(client.session)
        
        tools = original_tools
        tools_str = ".n".join([f"- {tool.name}: {tool.description}" for tool in tools])
        print(f"Available tools:\n{tools_str}")
        
        # Create LLM
        llm = ChatOpenAI(
            model="gpt-4o-mini",
            api_key="sk-proj-1234567890",
            temperature=0
        )
        
        # Simple system prompt - let LLM choose tools naturally
        system_prompt = f"""You are a helpful assistant that can access BuboBot monitoring data.

Available tools:
{tools_str}

Answer the user's question by using the appropriate tools to get current data.
"""

        # Create a ReAct agent that can reason about tool usage
        agent = create_react_agent(llm, tools, state_modifier=system_prompt, debug=True)
        
        # Run the agent with the user's query
        response = await agent.ainvoke({"messages": [("user", text)]})
        
        # Extract the final response
        final_message = response["messages"][-1]
        print(f"Agent response: {final_message.content}")
        return final_message.content
    finally:
        await client.close()

def main():
    """Interactive client for testing MCP server"""
    print("🤖 BuboBot MCP Client Interface")
    
    while True:
        text = input("Enter your question (or 'back' to change mode): ")
        if text.lower() == 'back':
            break
        # try:
        result = asyncio.run(query_with_running_server(text))
        print(f"\n✅ AI Response:\n{result}\n")
        # except Exception as e:
        #     print(f"❌ Error: {e}")

if __name__ == "__main__":
    main()
