"""
AgenticWerx MCP Client - Simple rule retrieval client

This module implements a simple MCP server that connects to your
AgenticWerx MCP server to retrieve rules.

Available Tools:
- get_rules: Get AgenticWerx marketplace rule packages
- get_custom_rules: Get user-uploaded custom rules (Pro/Pro Plus)
- get_team_rules: Get team-shared custom rules (Pro Plus teams)
- analyze_code: Analyze code with selected rules
- list_external_mcp_servers: Discover available external MCP servers (from Lambda)
- get_external_mcp_server: Get connection details for external servers (from Lambda)
- connect_external_server: Connect to external MCP servers (supports serverId or manual)
- list_external_connections: List connected external servers
- disconnect_external_server: Disconnect from external server
- call_external_tool: Call a tool on a connected external server
"""

import json
import logging
import os
from typing import Any

from mcp.server import Server
from mcp.server.models import InitializationOptions
from mcp.types import Resource, TextContent, Tool

from .api import AgenticWerxAPI, AgenticWerxAPIError
from .mcp_connection_manager import MCPConnectionManager

logger = logging.getLogger(__name__)


class AgenticWerxMCPClient:
    """
    Simple AgenticWerx MCP Client for rule retrieval.

    This client connects to your AgenticWerx MCP server and provides
    a simple interface to get rules through the MCP protocol.
    """

    def __init__(self, api_key: str, debug: bool = False):
        """
        Initialize the MCP client.

        Args:
            api_key: AgenticWerx API key
            debug: Enable debug logging
        """
        self.api_key = api_key
        self.debug = debug

        # Initialize the API client
        self.api = AgenticWerxAPI(api_key)

        # Initialize the MCP server
        self.server = Server("agenticwerx")

        # Initialize connection manager for external MCP servers with persistence
        connection_options = {
            "autoReconnect": True,  # Automatically restore connections on startup
        }
        # Only set custom persistence path if environment variable is set
        persistence_path = os.getenv("MCP_CONNECTIONS_FILE")
        if persistence_path:
            connection_options["persistencePath"] = persistence_path
        
        self.connection_manager = MCPConnectionManager(connection_options)

        # Cache for Lambda MCP server tools (fetched on startup)
        self.lambda_tools: list[dict[str, Any]] = []

        # Configure logging
        if debug:
            logging.getLogger().setLevel(logging.DEBUG)
        
        # Suppress httpx logging to avoid exposing URLs
        logging.getLogger("httpx").setLevel(logging.WARNING)

        logger.info("Initializing AgenticWerx MCP Client")

        # Set up MCP handlers
        self._setup_handlers()

    def _setup_handlers(self) -> None:
        """Set up MCP server handlers."""

        @self.server.list_resources()
        async def list_resources() -> list[Resource]:
            """List available rule resources."""
            logger.debug("Listing available rule resources")

            try:
                # Test if we can get rules from the Lambda MCP server
                await self.api.get_rules()

                # Create a single resource for all rules
                resource = Resource(
                    uri="agenticwerx://rules",
                    name="AgenticWerx Rules",
                    description="All available AgenticWerx rules from Lambda MCP server",
                    mimeType="application/json",
                )

                logger.info("Listed rule resources from Lambda MCP server")
                return [resource]

            except AgenticWerxAPIError as e:
                logger.error(f"API error listing resources: {e}")
                return []
            except Exception as e:
                logger.error(f"Unexpected error listing resources: {e}")
                return []

        @self.server.read_resource()
        async def read_resource(uri: str) -> str:
            """Read rule resource."""
            logger.debug(f"Reading resource: {uri}")

            if uri != "agenticwerx://rules":
                raise ValueError(f"Unknown resource URI: {uri}")

            try:
                rules_data = await self.api.get_rules()
                logger.debug("Successfully read rules resource from Lambda MCP server")
                return json.dumps(rules_data, indent=2)

            except AgenticWerxAPIError as e:
                logger.error(f"API error reading resource: {e}")
                raise
            except Exception as e:
                logger.error(f"Unexpected error reading resource: {e}")
                raise ValueError(f"Failed to read resource: {str(e)}") from e

        @self.server.list_tools()
        async def list_tools() -> list[Tool]:
            """List available tools."""
            logger.debug("Listing available tools")

            tools = []

            # Add tools from Lambda MCP server (dynamically fetched)
            for lambda_tool in self.lambda_tools:
                tools.append(
                    Tool(
                        name=lambda_tool["name"],
                        description=lambda_tool.get("description", ""),
                        inputSchema=lambda_tool.get("inputSchema", {"type": "object"}),
                    )
                )

            # Add client-side external MCP server management tools
            client_tools = [
                Tool(
                    name="connect_external_server",
                    description=(
                        "Connect to an external MCP server (AWS Documentation, GitHub, PostgreSQL, etc.). "
                        "Connections are automatically saved and restored on restart.\n\n"
                        "You can either:\n"
                        "1. Use list_external_mcp_servers and get_external_mcp_server to discover servers, then connect by serverId\n"
                        "2. Manually specify the command, args, and env\n\n"
                        "Example servers:\n"
                        "  • github: GitHub repository access\n"
                        "  • aws-documentation: AWS documentation search\n"
                        "  • postgres: PostgreSQL database access\n\n"
                        "Note: You'll need appropriate credentials in environment variables."
                    ),
                    inputSchema={
                        "type": "object",
                        "properties": {
                            "serverId": {
                                "type": "string",
                                "description": "Server ID from registry (e.g., 'github'). Use list_external_mcp_servers to discover. If provided, command/args are fetched automatically.",
                            },
                            "name": {
                                "type": "string",
                                "description": "Server name (optional if serverId is provided, otherwise required)",
                            },
                            "command": {
                                "type": "string",
                                "description": "Command to run (e.g., 'npx', 'uvx'). Required if serverId not provided.",
                            },
                            "args": {
                                "type": "array",
                                "items": {"type": "string"},
                                "description": "Command arguments. Required if serverId not provided.",
                            },
                            "env": {
                                "type": "object",
                                "description": "Environment variables required by the server",
                                "additionalProperties": {"type": "string"},
                            },
                        },
                        "additionalProperties": False,
                    },
                ),
                Tool(
                    name="list_external_connections",
                    description="List all currently connected external MCP servers and their available tools",
                    inputSchema={
                        "type": "object",
                        "properties": {},
                        "additionalProperties": False,
                    },
                ),
                Tool(
                    name="disconnect_external_server",
                    description="Disconnect from an external MCP server and remove from saved connections",
                    inputSchema={
                        "type": "object",
                        "properties": {
                            "serverName": {
                                "type": "string",
                                "description": "Name of the server to disconnect",
                            }
                        },
                        "required": ["serverName"],
                        "additionalProperties": False,
                    },
                ),
                Tool(
                    name="call_external_tool",
                    description=(
                        "Call a tool on a connected external MCP server. "
                        "This is the universal way to use any external server tool."
                    ),
                    inputSchema={
                        "type": "object",
                        "properties": {
                            "serverName": {
                                "type": "string",
                                "description": "Name of the connected server (e.g., 'github')",
                            },
                            "toolName": {
                                "type": "string",
                                "description": "Name of the tool to call (e.g., 'search_repositories')",
                            },
                            "arguments": {
                                "type": "object",
                                "description": "Arguments to pass to the tool",
                                "additionalProperties": True,
                            },
                        },
                        "required": ["serverName", "toolName"],
                        "additionalProperties": False,
                    },
                ),
            ]
            tools.extend(client_tools)

            # Add tools from connected external servers (with prefix)
            connections = self.connection_manager.list_connections()
            for conn in connections:
                connection = self.connection_manager.get_connection(conn["name"])
                if connection and connection.tools:
                    for tool in connection.tools:
                        tools.append(
                            Tool(
                                name=f"{conn['name']}__{tool['name']}",
                                description=f"[{conn['name']}] {tool.get('description', '')}",
                                inputSchema=tool.get("inputSchema", {"type": "object"}),
                            )
                        )

            logger.info(f"Listed {len(tools)} available tools ({len(self.lambda_tools)} from Lambda, {len(client_tools)} client-side, {len(tools) - len(self.lambda_tools) - len(client_tools)} from external connections)")
            return tools

        @self.server.call_tool()
        async def call_tool(name: str, arguments: dict[str, Any]) -> list[TextContent]:
            """Execute a tool."""
            logger.debug(f"Executing tool: {name}")

            try:
                # Check if this is an external server tool (with prefix)
                if "__" in name:
                    server_name, tool_name = name.split("__", 1)

                    logger.info(
                        f"Routing to external server: {server_name}.{tool_name}"
                    )

                    try:
                        result = await self.connection_manager.call_tool(
                            server_name, tool_name, arguments
                        )

                        logger.debug(
                            f"External tool call completed: {server_name}.{tool_name}"
                        )

                        # Return the result from the external server
                        if isinstance(result, dict) and "content" in result:
                            content = result["content"]
                            return [
                                TextContent(
                                    type=item.get("type", "text"),
                                    text=item.get("text", ""),
                                )
                                for item in content
                            ]
                        else:
                            return [
                                TextContent(
                                    type="text", text=json.dumps(result, indent=2)
                                )
                            ]

                    except Exception as e:
                        error_msg = f"Failed to call {server_name}.{tool_name}: {str(e)}"
                        logger.error(error_msg)
                        return [
                            TextContent(
                                type="text",
                                text=json.dumps({"error": error_msg}, indent=2),
                            )
                        ]
                # Otherwise, check if it's a Lambda MCP server tool
                elif any(t["name"] == name for t in self.lambda_tools):
                    logger.info(f"Routing to Lambda MCP server: {name}")
                    
                    try:
                        result = await self.api.call_tool(name, arguments)
                        
                        # Return the Lambda response directly
                        if isinstance(result, dict) and "content" in result:
                            content = result["content"]
                            return [
                                TextContent(
                                    type=item.get("type", "text"),
                                    text=item.get("text", ""),
                                )
                                for item in content
                            ]
                        else:
                            return [
                                TextContent(
                                    type="text", text=json.dumps(result, indent=2)
                                )
                            ]
                    
                    except Exception as e:
                        error_msg = f"Failed to call Lambda tool {name}: {str(e)}"
                        logger.error(error_msg)
                        return [
                            TextContent(
                                type="text",
                                text=json.dumps({"error": error_msg}, indent=2),
                            )
                        ]

                # Handle client-side external connection management tools
                if name == "connect_external_server":
                    # Handle connect_external_server tool
                    # Support two modes:
                    # 1. serverId (fetch config from Lambda)
                    # 2. Manual (name, command, args)

                    server_id = arguments.get("serverId")
                    server_name = arguments.get("name")
                    command = arguments.get("command")
                    args_list = arguments.get("args", [])
                    env_vars = arguments.get("env", {})

                    # Mode 1: serverId provided - fetch from Lambda
                    if server_id:
                        logger.info(
                            f"Connecting to external server via serverId: {server_id}"
                        )

                        try:
                            # Get server details from Lambda
                            server_details = await self.api.call_tool(
                                "get_external_mcp_server", {"serverId": server_id}
                            )

                            # Parse the server details from the text response
                            server_text = server_details["content"][0]["text"]

                            # Extract command and args from the response
                            import re

                            command_match = re.search(
                                r"\*\*Command:\*\* `([^`]+)`", server_text
                            )
                            args_match = re.search(
                                r"\*\*Arguments:\*\* (\[.*?\])", server_text
                            )

                            if not command_match or not args_match:
                                raise Exception("Could not parse server configuration")

                            command = command_match.group(1)
                            args_list = json.loads(args_match.group(1))

                            # Use serverId as the name if not provided
                            server_name = server_name or server_id

                            logger.info(
                                f"Fetched server config: {command} {' '.join(args_list)}"
                            )

                        except Exception as e:
                            error_msg = f"Failed to fetch server config for {server_id}: {str(e)}"
                            logger.error(error_msg)
                            return [
                                TextContent(
                                    type="text",
                                    text=json.dumps({"error": error_msg}, indent=2),
                                )
                            ]

                    # Mode 2: Manual connection
                    else:
                        if not server_name or not command or not args_list:
                            error_msg = "Missing required parameters: either serverId OR (name, command, args)"
                            logger.warning(error_msg)
                            return [
                                TextContent(
                                    type="text",
                                    text=json.dumps({"error": error_msg}, indent=2),
                                )
                            ]

                        logger.info(f"Connecting to external server manually: {server_name}")

                    # Connect to the external server
                    try:
                        result = await self.connection_manager.connect(
                            {
                                "name": server_name,
                                "command": command,
                                "args": args_list,
                                "env": env_vars,
                            }
                        )

                        response_text = (
                            f"✅ **Successfully Connected to {result['serverName']}**\n\n"
                            f"**Server Info:**\n"
                            f"- Name: {result['serverInfo']['name']}\n"
                            f"- Version: {result['serverInfo']['version']}\n\n"
                            f"**Available Tools ({len(result['tools'])}):**\n"
                        )

                        for tool in result["tools"]:
                            response_text += f"- `{tool['name']}`: {tool.get('description', 'No description')}\n"

                        response_text += (
                            f"\n**Usage:** Tools are now available with prefix `{result['serverName']}__`\n"
                            f"Example: `{result['serverName']}__{result['tools'][0]['name'] if result['tools'] else 'tool_name'}`\n\n"
                            f"**Note:** This connection will be automatically restored on restart."
                        )

                        logger.info(
                            f"Successfully connected to {server_name} with {len(result['tools'])} tools"
                        )

                        return [TextContent(type="text", text=response_text)]

                    except Exception as e:
                        error_msg = f"Failed to connect to {server_name}: {str(e)}"
                        logger.error(error_msg)
                        
                        # Enhanced error response with environment variable hints
                        error_response = {
                            "error": error_msg,
                            "server": server_name,
                        }
                        
                        # If we have a serverId, try to fetch required env vars from registry
                        if server_id:
                            try:
                                # Get server details from Lambda to extract required env vars
                                server_details = await self.api.call_tool(
                                    "get_external_mcp_server", {"serverId": server_id}
                                )
                                
                                # Parse the server details to extract env var names
                                server_text = server_details["content"][0]["text"]
                                
                                # Extract environment variables section
                                env_section_match = re.search(
                                    r"\*\*Environment Variables:\*\*(.*?)(?:\n\n|\Z)",
                                    server_text,
                                    re.DOTALL
                                )
                                
                                if env_section_match:
                                    env_section = env_section_match.group(1)
                                    # Extract variable names (format: - `VAR_NAME`: description)
                                    env_vars_found = re.findall(r"`([A-Z_]+)`", env_section)
                                    
                                    if env_vars_found:
                                        error_response["likely_cause"] = "Missing required environment variables"
                                        error_response["required_env_vars"] = env_vars_found
                                        error_response["help"] = f"Run get_external_mcp_server('{server_id}') for complete configuration details and examples."
                                
                            except Exception as fetch_error:
                                logger.debug(f"Could not fetch env var details: {fetch_error}")
                        
                        # Check for common timeout/connection errors
                        error_str = str(e).lower()
                        if "timeout" in error_str or "timed out" in error_str:
                            error_response["likely_cause"] = error_response.get(
                                "likely_cause", 
                                "Connection timeout - server may be unreachable or missing required configuration"
                            )
                        elif "connection" in error_str or "connect" in error_str:
                            error_response["likely_cause"] = error_response.get(
                                "likely_cause",
                                "Connection failed - check that the server command is installed and configured correctly"
                            )
                        
                        return [
                            TextContent(
                                type="text",
                                text=json.dumps(error_response, indent=2),
                            )
                        ]

                elif name == "list_external_connections":
                    # Handle list_external_connections tool
                    connections = self.connection_manager.list_connections()

                    response_text = "# Connected External MCP Servers\n\n"

                    if not connections:
                        response_text += "No external servers connected.\n\n"
                        response_text += "**To connect:** Use `connect_external_server` tool\n"
                    else:
                        response_text += f"**Total Connections:** {len(connections)}\n\n"
                        for conn in connections:
                            response_text += f"## {conn['name']}\n"
                            response_text += f"- **Status:** {'✅ Connected' if conn['isConnected'] else '❌ Disconnected'}\n"
                            response_text += f"- **Server:** {conn['serverInfo']['name']} v{conn['serverInfo']['version']}\n"
                            response_text += f"- **Tools:** {len(conn['tools'])} available\n"
                            response_text += f"  - {', '.join(conn['tools'])}\n\n"

                    logger.debug(f"Listed {len(connections)} external connections")
                    return [TextContent(type="text", text=response_text)]

                elif name == "disconnect_external_server":
                    # Handle disconnect_external_server tool
                    server_name = arguments.get("serverName")

                    if not server_name:
                        error_msg = "Missing required parameter: serverName"
                        logger.warning(error_msg)
                        return [
                            TextContent(
                                type="text",
                                text=json.dumps({"error": error_msg}, indent=2),
                            )
                        ]

                    logger.info(f"Disconnecting from external server: {server_name}")

                    result = await self.connection_manager.disconnect(server_name)

                    if result["success"]:
                        response_text = f"✅ Disconnected from {server_name}"
                        logger.info(f"Successfully disconnected from {server_name}")
                    else:
                        response_text = f"❌ {result.get('error', 'Unknown error')}"
                        logger.warning(f"Failed to disconnect from {server_name}")

                    return [TextContent(type="text", text=response_text)]

                elif name == "call_external_tool":
                    # Handle call_external_tool tool
                    server_name = arguments.get("serverName")
                    tool_name = arguments.get("toolName")
                    tool_args = arguments.get("arguments", {})

                    if not server_name or not tool_name:
                        error_msg = "Missing required parameters: serverName, toolName"
                        logger.warning(error_msg)
                        return [
                            TextContent(
                                type="text",
                                text=json.dumps({"error": error_msg}, indent=2),
                            )
                        ]

                    logger.info(
                        f"Calling external tool: {server_name}.{tool_name}"
                    )

                    try:
                        result = await self.connection_manager.call_tool(
                            server_name, tool_name, tool_args
                        )

                        logger.debug(
                            f"External tool call completed: {server_name}.{tool_name}"
                        )

                        # Return the result from the external server
                        # The result should already be in MCP format
                        if isinstance(result, dict) and "content" in result:
                            # Extract content array
                            content = result["content"]
                            return [
                                TextContent(
                                    type=item.get("type", "text"),
                                    text=item.get("text", ""),
                                )
                                for item in content
                            ]
                        else:
                            # Fallback: wrap in JSON
                            return [
                                TextContent(
                                    type="text", text=json.dumps(result, indent=2)
                                )
                            ]

                    except Exception as e:
                        error_msg = f"Failed to call {server_name}.{tool_name}: {str(e)}"
                        logger.error(error_msg)
                        return [
                            TextContent(
                                type="text",
                                text=json.dumps({"error": error_msg}, indent=2),
                            )
                        ]

                # Check if this is a call to an external server tool (with prefix)
                elif "__" in name:
                    server_name, tool_name = name.split("__", 1)

                    logger.info(
                        f"Routing to external server: {server_name}.{tool_name}"
                    )

                    try:
                        result = await self.connection_manager.call_tool(
                            server_name, tool_name, arguments
                        )

                        logger.debug(
                            f"External tool call completed: {server_name}.{tool_name}"
                        )

                        # Return the result from the external server
                        if isinstance(result, dict) and "content" in result:
                            content = result["content"]
                            return [
                                TextContent(
                                    type=item.get("type", "text"),
                                    text=item.get("text", ""),
                                )
                                for item in content
                            ]
                        else:
                            return [
                                TextContent(
                                    type="text", text=json.dumps(result, indent=2)
                                )
                            ]

                    except Exception as e:
                        error_msg = f"Failed to call {server_name}.{tool_name}: {str(e)}"
                        logger.error(error_msg)
                        return [
                            TextContent(
                                type="text",
                                text=json.dumps({"error": error_msg}, indent=2),
                            )
                        ]

                else:
                    # Unsupported tool
                    error_msg = f"Tool '{name}' is not supported. This tool may not be available from the Lambda MCP server or external connections."
                    logger.warning(f"Unsupported tool requested: {name}")
                    return [
                        TextContent(
                            type="text", text=json.dumps({"error": error_msg}, indent=2)
                        )
                    ]

            except AgenticWerxAPIError as e:
                error_msg = f"AgenticWerx API Error: {str(e)}"
                logger.error(f"API error in tool {name}: {e}")
                return [
                    TextContent(
                        type="text", text=json.dumps({"error": error_msg}, indent=2)
                    )
                ]

            except Exception as e:
                error_msg = f"Tool execution error: {str(e)}"
                logger.error(f"Unexpected error in tool {name}: {e}")
                return [
                    TextContent(
                        type="text", text=json.dumps({"error": error_msg}, indent=2)
                    )
                ]

    async def _fetch_lambda_tools(self) -> None:
        """Fetch available tools from the Lambda MCP server."""
        try:
            logger.info("Fetching tools from Lambda MCP server...")
            self.lambda_tools = await self.api.list_tools()
            logger.info(f"Loaded {len(self.lambda_tools)} tools from Lambda MCP server")
            
            # Log tool names for debugging
            tool_names = [t["name"] for t in self.lambda_tools]
            logger.debug(f"Lambda tools: {', '.join(tool_names)}")
        except Exception as e:
            logger.error(f"Failed to fetch Lambda tools: {e}")
            logger.warning("Continuing with empty Lambda tool list")
            self.lambda_tools = []

    async def test_connection(self) -> bool:
        """Test connection to the Lambda MCP server."""
        return await self.api.test_connection()

    async def run(self) -> None:
        """Run the MCP server."""
        logger.info("Starting AgenticWerx MCP Client")

        # Initialize connection manager and restore saved connections
        try:
            await self.connection_manager.initialize()
            logger.info(
                f"Connection manager initialized (persistence: {self.connection_manager.get_persistence_path()})"
            )
        except Exception as e:
            logger.error(f"Failed to initialize connection manager: {e}")

        # Test connection to Lambda MCP server on startup
        logger.info("Testing connection to Lambda MCP server...")
        connection_ok = await self.test_connection()
        if not connection_ok:
            logger.error("Failed to connect to Lambda MCP server")
            # Continue anyway - the client might still work for some operations
        else:
            logger.info("Successfully connected to Lambda MCP server")
        
        # Fetch available tools from Lambda MCP server
        await self._fetch_lambda_tools()

        try:
            from mcp.server.stdio import stdio_server

            async with stdio_server() as (read_stream, write_stream):
                logger.info("MCP Client started, waiting for connections...")

                from mcp.types import (
                    ResourcesCapability,
                    ServerCapabilities,
                    ToolsCapability,
                )

                init_options = InitializationOptions(
                    server_name="agenticwerx",
                    server_version="1.0.0",
                    capabilities=ServerCapabilities(
                        resources=ResourcesCapability(
                            subscribe=False, listChanged=False
                        ),
                        tools=ToolsCapability(listChanged=False),
                    ),
                )

                await self.server.run(read_stream, write_stream, init_options)

        except Exception as e:
            logger.error(f"Error running MCP server: {e}")
            raise
        finally:
            # Disconnect all external servers
            await self.connection_manager.disconnect_all()
            await self.api.close()
            logger.info("AgenticWerx MCP Client stopped")

    async def __aenter__(self):
        """Async context manager entry."""
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        """Async context manager exit."""
        await self.api.close()
