"""
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:
- All tools from the AgenticWerx Lambda MCP server are dynamically loaded
"""

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"}),
                    )
                )

            logger.info(f"Listed {len(tools)} available tools from Lambda MCP server")
            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 it's a Lambda MCP server tool
                if 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),
                            )
                        ]

                else:
                    # Unsupported tool
                    error_msg = f"Tool '{name}' is not supported. This tool may not be available from the Lambda MCP server."
                    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()
