"""
Copyright (C) 2025 Luciano Guerche

This file is part of rabbitmq-mcp-server.

rabbitmq-mcp-server is free software: you can redistribute it and/or modify
it under the terms of the GNU Lesser General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

rabbitmq-mcp-server is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
GNU Lesser General Public License for more details.

You should have received a copy of the GNU Lesser General Public License
along with rabbitmq-mcp-server. If not, see <https://www.gnu.org/licenses/>.
"""

"""MCP Server implementation with JSON-RPC 2.0 over stdio transport."""

import asyncio
import json
import signal
import sys
import time
from logging.correlation import generate_correlation_id
from pathlib import Path
from typing import Any

from mcp.server import Server
from mcp.server.stdio import stdio_server

from config import load_settings
from config.logging import get_logger
from config.settings import Settings
from rabbitmq_mcp_connection.connection.health import HealthChecker
from rabbitmq_mcp_connection.connection.manager import ConnectionManager
from rabbitmq_mcp_connection.executor import OperationExecutor
from rabbitmq_mcp_connection.http_client import HTTPClient, HTTPClientManager
from rabbitmq_mcp_connection.reconnect import ReconnectionManager
from rabbitmq_mcp_connection.schemas.connection import ConnectionConfig
from rabbitmq_mcp_connection.tools.call_id_http import handle_call_id_http
from rabbitmq_mcp_connection.tools.get_id import handle_get_id
from tools.call_id import call_id
from tools.search_ids import SearchResult, search_ids

logger = get_logger(__name__)

# Server info constants
SERVER_NAME = "rabbitmq-mcp-server"
SERVER_VERSION = "3.0.0"
PROTOCOL_VERSION = "2024-11-05"

# Graceful shutdown parameters
SHUTDOWN_GRACE_PERIOD = 5  # seconds


class MCPServer:
    """MCP Server with JSON-RPC 2.0 protocol over stdio transport."""

    def __init__(self, settings: Settings | None = None) -> None:
        """Initialize MCP server with configuration."""
        self.settings = settings or Settings()
        self.server = Server(SERVER_NAME)
        self.shutdown_flag = False
        self.operations_registry: dict[str, Any] = {}
        self.embeddings: dict[str, Any] = {}
        self.http_client: HTTPClient | None = None
        self.http_client_manager: HTTPClientManager | None = None  # Story 2.3
        self.operation_executor: OperationExecutor | None = None
        self.connection_manager: ConnectionManager | None = None
        self.health_checker: HealthChecker | None = None  # Story 2.4
        self.reconnection_manager: ReconnectionManager | None = None  # Story 2.5
        self._register_signal_handlers()
        self._register_tools()

    def _register_signal_handlers(self) -> None:
        """Register handlers for SIGTERM and SIGINT."""
        signal.signal(signal.SIGTERM, self._handle_shutdown_signal)
        signal.signal(signal.SIGINT, self._handle_shutdown_signal)

    def _handle_shutdown_signal(self, signum: int, frame: Any) -> None:
        """Handle shutdown signals."""
        logger.info("shutdown_signal_received", signal=signum)
        self.shutdown_flag = True

    def _load_operations_registry(self) -> None:
        """Load operations.json into memory."""
        start_time = time.time()
        operations_path = Path(self.settings.operations_registry_path)

        if not operations_path.exists():
            error_msg = f"Operations registry not found at {operations_path}"
            logger.error(
                "operations_registry_load_failed",
                error=error_msg,
                path=str(operations_path),
            )
            raise FileNotFoundError(error_msg)

        try:
            with operations_path.open("r") as f:
                self.operations_registry = json.load(f)

            # Validate version consistency
            registry_version = self.operations_registry.get("api_version")
            if registry_version != self.settings.api_version:
                error_msg = (
                    f"API version mismatch: operations registry has version '{registry_version}' "
                    f"but settings expect '{self.settings.api_version}'"
                )
                logger.error("operations_version_mismatch", error=error_msg)
                raise ValueError(error_msg)

            load_time = (time.time() - start_time) * 1000
            logger.info(
                "operations_registry_loaded",
                path=str(operations_path),
                api_version=self.settings.api_version,
                load_time_ms=f"{load_time:.2f}",
                operation_count=len(self.operations_registry.get("operations", {})),
            )
        except json.JSONDecodeError as e:
            logger.error(
                "operations_registry_invalid_json",
                error=str(e),
                path=str(operations_path),
            )
            raise ValueError(f"Invalid JSON in operations registry: {e}")

    def _load_embeddings(self) -> None:
        """Load embeddings.json into memory."""
        start_time = time.time()
        embeddings_path = Path(self.settings.embeddings_path)

        if not embeddings_path.exists():
            error_msg = f"Embeddings file not found at {embeddings_path}"
            logger.error(
                "embeddings_load_failed", error=error_msg, path=str(embeddings_path)
            )
            raise FileNotFoundError(error_msg)

        try:
            with embeddings_path.open("r") as f:
                self.embeddings = json.load(f)

            # Validate version consistency
            embeddings_version = self.embeddings.get("api_version")
            if embeddings_version != self.settings.api_version:
                error_msg = (
                    f"API version mismatch: embeddings file has version '{embeddings_version}' "
                    f"but settings expect '{self.settings.api_version}'"
                )
                logger.error("embeddings_version_mismatch", error=error_msg)
                raise ValueError(error_msg)

            load_time = (time.time() - start_time) * 1000
            logger.info(
                "embeddings_loaded",
                path=str(embeddings_path),
                api_version=self.settings.api_version,
                load_time_ms=f"{load_time:.2f}",
                embedding_count=len(self.embeddings.get("embeddings", {})),
            )
        except json.JSONDecodeError as e:
            logger.error(
                "embeddings_invalid_json", error=str(e), path=str(embeddings_path)
            )
            raise ValueError(f"Invalid JSON in embeddings file: {e}")

    def _register_tools(self) -> None:
        """Register MCP tools with the server."""

        @self.server.get_server_info()  # type: ignore[misc]
        async def get_server_info() -> dict[str, Any]:
            """Return server information including API version."""
            return {
                "name": SERVER_NAME,
                "version": SERVER_VERSION,
                "protocol_version": PROTOCOL_VERSION,
                "api_version": self.settings.api_version,
            }

        @self.server.list_tools()  # type: ignore[misc]
        async def list_tools() -> list[dict[str, Any]]:
            """Return list of available MCP tools."""
            return [
                {
                    "name": "search-ids",
                    "description": "Search for RabbitMQ operation IDs using semantic search",
                    "inputSchema": {
                        "type": "object",
                        "properties": {
                            "query": {
                                "type": "string",
                                "description": "Natural language query to search for operations",
                            },
                            "threshold": {
                                "type": "number",
                                "description": "Minimum similarity score (0.0-1.0, default: 0.7)",
                                "default": 0.7,
                                "minimum": 0.0,
                                "maximum": 1.0,
                            },
                            "max_results": {
                                "type": "integer",
                                "description": "Maximum results to return (default: 10, max: 50)",
                                "default": 10,
                                "minimum": 1,
                                "maximum": 50,
                            },
                        },
                        "required": ["query"],
                    },
                },
                {
                    "name": "get-id",
                    "description": "Get detailed operation schema and documentation by operation ID",
                    "inputSchema": {
                        "type": "object",
                        "properties": {
                            "operation_id": {
                                "type": "string",
                                "description": "The operation ID to retrieve",
                            }
                        },
                        "required": ["operation_id"],
                    },
                },
                {
                    "name": "call-id",
                    "description": "Execute a RabbitMQ operation by operation ID with parameters",
                    "inputSchema": {
                        "type": "object",
                        "properties": {
                            "operation_id": {
                                "type": "string",
                                "description": "The operation ID to execute",
                            },
                            "parameters": {
                                "type": "object",
                                "description": "Operation-specific parameters",
                                "default": {},
                            },
                        },
                        "required": ["operation_id"],
                    },
                },
                {
                    "name": "connection.health",
                    "description": "Get current RabbitMQ connection health status (AMQP and HTTP)",
                    "inputSchema": {
                        "type": "object",
                        "properties": {},
                    },
                },
                {
                    "name": "connection.reconnection_state",
                    "description": "Get current reconnection state including attempt number, delay, and metrics",
                    "inputSchema": {
                        "type": "object",
                        "properties": {},
                    },
                },
                {
                    "name": "connection.pool_status",
                    "description": "Get HTTP connection pool health metrics including active connections, pool size, health score, and eviction statistics",
                    "inputSchema": {
                        "type": "object",
                        "properties": {},
                    },
                },
            ]

        @self.server.call_tool()  # type: ignore[misc]
        async def call_tool(
            name: str, arguments: dict[str, Any]
        ) -> list[dict[str, Any]]:
            """Route tool calls to appropriate implementation."""
            correlation_id = generate_correlation_id()

            try:
                if name == "search-ids":
                    search_result: SearchResult = search_ids(
                        query=arguments["query"],
                        threshold=arguments.get("threshold", 0.7),
                        max_results=arguments.get("max_results", 10),
                        embeddings=self.embeddings,
                        operations_registry=self.operations_registry,
                        correlation_id=correlation_id,
                    )
                    return [{"type": "text", "text": search_result.model_dump_json()}]

                elif name == "get-id":
                    # Use new handle_get_id implementation from rabbitmq_mcp_connection
                    result = await handle_get_id(
                        operation_id=arguments["operation_id"],
                        correlation_id=correlation_id,
                    )

                    # Handle MCPToolResult format
                    if result.success and result.result:
                        # Return formatted schema as text
                        return [{"type": "text", "text": result.result["formatted"]}]
                    elif result.error:
                        # Return error in MCP format
                        error_response: dict[str, Any] = {
                            "error": result.error.code,
                            "message": result.error.message,
                        }
                        if result.error.details:
                            error_response["details"] = result.error.details
                        return [{"type": "text", "text": json.dumps(error_response)}]
                    else:
                        # Fallback for unexpected case
                        return [
                            {
                                "type": "text",
                                "text": json.dumps(
                                    {
                                        "error": "UNKNOWN_ERROR",
                                        "message": "Unexpected result format",
                                    }
                                ),
                            }
                        ]

                elif name == "call-id":
                    # Check if operation is HTTP-based or connection-based
                    operation_id = arguments["operation_id"]

                    # HTTP-based operations from operations.json (most operations)
                    if (
                        self.operation_executor
                        and operation_id
                        in self.operations_registry.get("operations", {})
                    ):
                        result = await handle_call_id_http(
                            operation_id=operation_id,
                            parameters=arguments.get("parameters", {}),
                            executor=self.operation_executor,
                        )

                        # Handle MCPToolResult format
                        if result.success and result.result:
                            # Return formatted result as text
                            formatted = result.result.get(
                                "formatted", json.dumps(result.result)
                            )
                            return [{"type": "text", "text": formatted}]
                        elif result.error:
                            # Return error in MCP format
                            error_response_call: dict[str, Any] = {
                                "error": result.error.code,
                                "message": result.error.message,
                            }
                            if result.error.details:
                                error_response_call["details"] = result.error.details
                            return [
                                {
                                    "type": "text",
                                    "text": json.dumps(error_response_call),
                                }
                            ]
                        else:
                            # Fallback for unexpected case
                            return [
                                {
                                    "type": "text",
                                    "text": json.dumps(
                                        {
                                            "error": "UNKNOWN_ERROR",
                                            "message": "Unexpected result format",
                                        }
                                    ),
                                }
                            ]

                    # Fallback to existing call_id handler for topology operations
                    call_result = call_id(
                        endpoint_id=operation_id,
                        params=arguments.get("parameters", {}),
                        pagination=None,
                    )
                    # call_id returns BaseModel instances or dict - handle both
                    if hasattr(call_result, "model_dump_json"):
                        result_json = call_result.model_dump_json()
                    else:
                        result_json = json.dumps(call_result)
                    return [{"type": "text", "text": result_json}]

                elif name == "connection.health":
                    # Get health status (Story 2.4 - AC #8)
                    if self.health_checker:
                        status = self.health_checker.get_status()
                        if status:
                            return [{"type": "text", "text": status.model_dump_json()}]
                        else:
                            return [
                                {
                                    "type": "text",
                                    "text": json.dumps(
                                        {
                                            "error": "HEALTH_NOT_AVAILABLE",
                                            "message": "Health check service not yet initialized",
                                        }
                                    ),
                                }
                            ]
                    else:
                        return [
                            {
                                "type": "text",
                                "text": json.dumps(
                                    {
                                        "error": "HEALTH_SERVICE_DISABLED",
                                        "message": "Health check service not configured",
                                    }
                                ),
                            }
                        ]

                elif name == "connection.reconnection_state":
                    # Get reconnection state (Story 2.5 - AC #8)
                    if self.reconnection_manager:
                        state = self.reconnection_manager.get_state()
                        return [{"type": "text", "text": state.model_dump_json()}]
                    else:
                        return [
                            {
                                "type": "text",
                                "text": json.dumps(
                                    {
                                        "error": "RECONNECTION_SERVICE_DISABLED",
                                        "message": "Reconnection manager not configured",
                                    }
                                ),
                            }
                        ]

                elif name == "connection.pool_status":
                    # Get HTTP connection pool status (Story 2.6 - AC #7)
                    if self.http_client:
                        pool_status = self.http_client.get_pool_status()
                        return [{"type": "text", "text": pool_status.model_dump_json()}]
                    else:
                        return [
                            {
                                "type": "text",
                                "text": json.dumps(
                                    {
                                        "error": "HTTP_CLIENT_DISABLED",
                                        "message": "HTTP client not configured",
                                    }
                                ),
                            }
                        ]

                else:
                    raise ValueError(f"Unknown tool: {name}")

            except ValueError as e:
                # Parameter validation errors - return as-is with correlation_id
                logger.error(
                    "tool_validation_error",
                    tool=name,
                    error=str(e),
                    correlation_id=correlation_id,
                )
                raise
            except Exception as e:
                logger.error(
                    "tool_execution_error",
                    tool=name,
                    error=str(e),
                    correlation_id=correlation_id,
                )
                raise ValueError(f"Tool execution failed: {str(e)}") from e

    async def startup(self) -> None:
        """Initialize server and load required artifacts."""
        startup_start = time.time()

        # Format version display message
        version_display = (
            f"{SERVER_NAME} v{SERVER_VERSION} (API: {self.settings.api_version}.x)"
        )
        logger.info(
            "server_starting",
            name=SERVER_NAME,
            version=SERVER_VERSION,
            api_version=self.settings.api_version,
            display=version_display,
        )
        print(
            f"RabbitMQ MCP Server v{SERVER_VERSION} (API: {self.settings.api_version}.x)"
        )

        try:
            # Load artifacts
            self._load_operations_registry()
            self._load_embeddings()

            # Initialize HTTP client for Management API operations (Story 1.10)
            self.http_client = HTTPClient(
                base_url=self.settings.rabbitmq_base_url,
                username=self.settings.rabbitmq_username,
                password=self.settings.rabbitmq_password.get_secret_value(),
                verify_ssl=self.settings.rabbitmq_verify_ssl,
                timeout=self.settings.rabbitmq_timeout,
                pool_size=self.settings.rabbitmq_pool_size,
            )

            # Initialize HTTP client manager with new Settings fields (Story 2.3)
            self.http_client_manager = HTTPClientManager(
                host=self.settings.http_host,
                port=self.settings.http_port,
                username=self.settings.http_user,
                password=self.settings.http_password.get_secret_value(),
                use_tls=self.settings.http_use_tls,
                verify_cert=self.settings.http_verify_cert,
                ca_bundle=self.settings.http_ca_bundle,
                timeout=self.settings.http_timeout,
                pool_size=self.settings.max_http_connections,
                pool_health_check_interval=self.settings.pool_health_check_interval,
            )

            # Test HTTP connection with GET /api/overview (Story 2.3 - AC #1)
            try:
                response = await self.http_client_manager.request("GET", "/overview")
                if response.status_code == 200:
                    logger.info(
                        "http_management_api_connected",
                        base_url=self.http_client_manager.base_url,
                        status_code=response.status_code,
                    )
                    print(
                        f"HTTP Management API client initialized successfully at {self.http_client_manager.base_url}"
                    )
            except Exception as e:
                logger.warning(
                    "http_management_api_connection_failed",
                    error=str(e),
                    base_url=self.http_client_manager.base_url,
                    message="HTTP-only operations may be unavailable",
                )
                # Don't prevent server startup - HTTP connection is optional

            # Initialize operation executor
            self.operation_executor = OperationExecutor(
                http_client=self.http_client,
                operation_registry=self.operations_registry,
            )

            logger.info(
                "http_client_initialized",
                base_url=self.settings.rabbitmq_base_url,
                pool_size=self.settings.rabbitmq_pool_size,
                timeout=self.settings.rabbitmq_timeout,
            )

            # Initialize AMQP connection manager (Story 2.2)
            amqp_config = ConnectionConfig(
                host=self.settings.amqp_host,
                port=self.settings.amqp_port,
                user=self.settings.amqp_user,
                password=self.settings.amqp_password.get_secret_value(),
                vhost=self.settings.amqp_vhost,
                use_tls=self.settings.amqp_use_tls,
                verify_cert=self.settings.amqp_verify_cert,
                ca_bundle=self.settings.amqp_ca_bundle,
                timeout=30,
                heartbeat=60,
            )
            self.connection_manager = ConnectionManager(amqp_config)

            try:
                await self.connection_manager.connect()
                logger.info(
                    "amqp_connection_established",
                    host=self.settings.amqp_host,
                    port=self.settings.amqp_port,
                    vhost=self.settings.amqp_vhost,
                )
            except Exception as e:
                logger.error(
                    "amqp_connection_failed",
                    error=str(e),
                    host=self.settings.amqp_host,
                    port=self.settings.amqp_port,
                )
                # Note: Not failing startup - AMQP connection is optional for HTTP-only operations
                logger.warning(
                    "server_starting_without_amqp",
                    message="Server will operate in HTTP-only mode",
                )

            # Initialize reconnection manager (Story 2.5 - AC #1-9)
            if self.connection_manager and self.http_client_manager:
                self.reconnection_manager = ReconnectionManager(
                    amqp_manager=self.connection_manager,
                    http_manager=self.http_client_manager,
                    max_delay=self.settings.reconnect_max_delay,
                    reset_after=self.settings.reconnect_reset_after,
                )
                logger.info(
                    "reconnection_manager_initialized",
                    max_delay=self.settings.reconnect_max_delay,
                    reset_after=self.settings.reconnect_reset_after,
                    msg=f"Reconnection manager initialized (max_delay: {self.settings.reconnect_max_delay}s, reset_after: {self.settings.reconnect_reset_after}s)",
                )

            # Initialize health check service (Story 2.4 - AC #4, #8)
            if self.connection_manager and self.http_client_manager:
                self.health_checker = HealthChecker(
                    manager=self.connection_manager,
                    timeout=float(self.settings.health_check_timeout),
                    http_client=self.http_client_manager,
                    check_interval=self.settings.health_check_interval,
                    reconnection_manager=self.reconnection_manager,
                )
                await self.health_checker.start()
                logger.info(
                    "health_check_service_started",
                    interval=self.settings.health_check_interval,
                    timeout=self.settings.health_check_timeout,
                    msg=f"Health check service started (interval: {self.settings.health_check_interval}s, timeout: {self.settings.health_check_timeout}s)",
                )

                # Start proactive pool health checks (Story 2.6 - AC #10)
                await self.http_client_manager.start_health_checks()
                logger.info(
                    "pool_health_checks_started",
                    interval=self.settings.pool_health_check_interval,
                    msg=f"HTTP pool health checks started (interval: {self.settings.pool_health_check_interval}s)",
                )

            startup_time = (time.time() - startup_start) * 1000

            if startup_time > 1000:
                logger.warning(
                    "server_startup_slow",
                    startup_time_ms=f"{startup_time:.2f}",
                    target_ms="1000",
                )

            logger.info(
                "server_ready",
                startup_time_ms=f"{startup_time:.2f}",
                protocol_version=PROTOCOL_VERSION,
                api_version=self.settings.api_version,
            )

        except Exception as e:
            logger.error("server_startup_failed", error=str(e))
            raise

    async def shutdown(self) -> None:
        """Perform graceful shutdown."""
        logger.info("shutdown_initiated")

        # Stop health check service (Story 2.4)
        if self.health_checker:
            await self.health_checker.stop()
            logger.info("health_check_service_stopped")

        # Stop reconnection manager (Story 2.5)
        if self.reconnection_manager:
            await self.reconnection_manager.shutdown()
            logger.info("reconnection_manager_stopped")

        # Close AMQP connection
        if self.connection_manager:
            await self.connection_manager.disconnect()
            logger.info("amqp_connection_closed")

        # Close HTTP client manager (Story 2.3)
        if self.http_client_manager:
            await self.http_client_manager.close()
            logger.info("http_client_manager_closed")

        # Close legacy HTTP client
        if self.http_client:
            await self.http_client.close()
            logger.info("http_client_closed")

        # Allow in-flight requests to complete
        try:
            await asyncio.wait_for(
                asyncio.sleep(0.1),  # Brief pause for request completion
                timeout=SHUTDOWN_GRACE_PERIOD,
            )
        except TimeoutError:
            logger.warning(
                "shutdown_grace_period_expired",
                grace_period_seconds=SHUTDOWN_GRACE_PERIOD,
            )

        logger.info("shutdown_complete")

    async def run(self) -> None:
        """Run the MCP server with stdio transport."""
        await self.startup()

        try:
            async with stdio_server() as (read_stream, write_stream):
                await self.server.run(
                    read_stream,
                    write_stream,
                    self.server.create_initialization_options(),
                )
        except Exception as e:
            logger.error("server_runtime_error", error=str(e))
            raise
        finally:
            if self.shutdown_flag:
                await self.shutdown()


async def main(cli_args: dict[str, Any] | None = None) -> None:
    """Main entry point for MCP server."""
    settings = load_settings(cli_args)
    server = MCPServer(settings)

    try:
        await server.run()
    except KeyboardInterrupt:
        logger.info("server_interrupted")
    except Exception as e:
        logger.error("server_fatal_error", error=str(e))
        sys.exit(1)

    sys.exit(0)


if __name__ == "__main__":
    asyncio.run(main())
