# -*- coding: utf-8 -*-
"""Location: ./mcpgateway/reverse_proxy.py
Copyright 2025
SPDX-License-Identifier: Apache-2.0
Authors: Mihai Criveti

MCP Reverse Proxy - Bridge local MCP servers to remote gateways.
This module implements a reverse proxy that connects local MCP servers
(running via stdio) to remote gateways, enabling servers behind firewalls
or NATs to be accessible without inbound network access.

The reverse proxy establishes an outbound WebSocket or SSE connection to
a remote gateway and registers the local server. All MCP protocol messages
are then tunneled through this persistent connection.

Environment variables:
- REVERSE_PROXY_GATEWAY: Remote gateway URL (required)
- REVERSE_PROXY_TOKEN: Bearer token for authentication (optional)
- REVERSE_PROXY_RECONNECT_DELAY: Initial reconnection delay in seconds (default 1)
- REVERSE_PROXY_MAX_RETRIES: Maximum reconnection attempts (default 0 = infinite)
- REVERSE_PROXY_LOG_LEVEL: Python log level (default INFO)

Example:
    $ export REVERSE_PROXY_GATEWAY=https://gateway.example.com
    $ export REVERSE_PROXY_TOKEN=$(python3 -m mcpgateway.utils.create_jwt_token --username admin@example.com --exp 10080 --secret key)
    $ python3 -m mcpgateway.reverse_proxy --local-stdio "uvx mcp-server-git"
"""

# Future
from __future__ import annotations

# Standard
import argparse
import asyncio
from contextlib import suppress
from enum import Enum
import json
import logging
import os
import shlex
import signal
import sys
from typing import Any, cast, Dict, List, Optional
from urllib.parse import urljoin, urlparse
import uuid

try:
    # Third-Party
    import httpx
except ImportError:
    httpx = None  # type: ignore[assignment]

try:
    # Third-Party
    import websockets
except ImportError:
    websockets = None  # type: ignore[assignment]


try:
    # Third-Party
    import yaml
except ImportError:
    yaml = None  # type: ignore[assignment]

# First-Party
from mcpgateway.services.logging_service import LoggingService

# Type alias for the websocket client protocol to avoid hard dependency at type-check time
WSClientProtocol = Any  # type: ignore[assignment]

# Initialize logging
logging_service = LoggingService()
LOGGER = logging_service.get_logger("mcpgateway.reverse_proxy")

# Environment variable names
ENV_GATEWAY = "REVERSE_PROXY_GATEWAY"
ENV_TOKEN = "REVERSE_PROXY_TOKEN"  # nosec B105 - environment variable name, not a secret
ENV_RECONNECT_DELAY = "REVERSE_PROXY_RECONNECT_DELAY"
ENV_MAX_RETRIES = "REVERSE_PROXY_MAX_RETRIES"
ENV_LOG_LEVEL = "REVERSE_PROXY_LOG_LEVEL"

# Default configuration
DEFAULT_RECONNECT_DELAY = 1.0  # seconds
DEFAULT_MAX_RETRIES = 0  # 0 = infinite
DEFAULT_KEEPALIVE_INTERVAL = 30  # seconds
DEFAULT_REQUEST_TIMEOUT = 90  # seconds


class ConnectionState(Enum):
    """Connection state enumeration.

    Examples:
        >>> ConnectionState.DISCONNECTED.value
        'disconnected'
        >>> ConnectionState.CONNECTED.value
        'connected'
        >>> ConnectionState.CONNECTING.value
        'connecting'
    """

    DISCONNECTED = "disconnected"
    CONNECTING = "connecting"
    CONNECTED = "connected"
    RECONNECTING = "reconnecting"
    SHUTTING_DOWN = "shutting_down"


class MessageType(Enum):
    """Control message types for the reverse proxy protocol.

    Examples:
        >>> MessageType.REGISTER.value
        'register'
        >>> MessageType.REQUEST.value
        'request'
        >>> MessageType.HEARTBEAT.value
        'heartbeat'
    """

    # Control messages
    REGISTER = "register"
    UNREGISTER = "unregister"
    HEARTBEAT = "heartbeat"
    ERROR = "error"

    # MCP messages
    REQUEST = "request"
    RESPONSE = "response"
    NOTIFICATION = "notification"


class StdioProcess:
    """Manages a local MCP server subprocess via stdio."""

    def __init__(self, command: str):
        """Initialize stdio process manager.

        Args:
            command: The command to run as a subprocess.
        """
        self.command = command
        self.process: Optional[asyncio.subprocess.Process] = None
        self._stdout_reader_task: Optional[asyncio.Task] = None
        self._message_handlers: List[Any] = []

    async def start(self) -> None:
        """Start the stdio subprocess.

        Raises:
            RuntimeError: If subprocess creation fails with stdio.
        """
        LOGGER.info(f"Starting local MCP server: {self.command}")

        self.process = await asyncio.create_subprocess_exec(
            *shlex.split(self.command),
            stdin=asyncio.subprocess.PIPE,
            stdout=asyncio.subprocess.PIPE,
            stderr=sys.stderr,  # Pass through for debugging
        )

        if not self.process.stdin or not self.process.stdout:
            raise RuntimeError(f"Failed to create subprocess with stdio: {self.command}")

        # Start reading stdout
        self._stdout_reader_task = asyncio.create_task(self._read_stdout())
        LOGGER.info(f"Local MCP server started (PID: {self.process.pid})")

    async def stop(self) -> None:
        """Stop the stdio subprocess gracefully."""
        if not self.process:
            return

        LOGGER.info(f"Stopping local MCP server (PID: {self.process.pid})")

        # Cancel stdout reader
        if self._stdout_reader_task:
            self._stdout_reader_task.cancel()
            with suppress(asyncio.CancelledError):
                await self._stdout_reader_task

        # Terminate process
        self.process.terminate()
        with suppress(asyncio.TimeoutError):
            await asyncio.wait_for(self.process.wait(), timeout=5)

        # Force kill if needed
        if self.process.returncode is None:
            LOGGER.warning("Force killing subprocess")
            self.process.kill()
            await self.process.wait()

    async def send(self, message: str) -> None:
        """Send a message to the subprocess stdin.

        Args:
            message: JSON-RPC message to send.

        Raises:
            RuntimeError: If subprocess is not running.
        """
        if not self.process or not self.process.stdin:
            raise RuntimeError("Subprocess not running")

        LOGGER.debug(f"→ stdio: {message[:200]}...")
        self.process.stdin.write((message + "\n").encode())
        await self.process.stdin.drain()

    def add_message_handler(self, handler) -> None:
        """Add a handler for messages from stdout.

        Args:
            handler: Async function to handle messages.
        """
        self._message_handlers.append(handler)

    async def _read_stdout(self) -> None:
        """Read messages from subprocess stdout.

        Raises:
            asyncio.CancelledError: When the read task is cancelled.
        """
        if not self.process or not self.process.stdout:
            return

        try:
            while True:
                line = await self.process.stdout.readline()
                if not line:
                    break

                message = line.decode().strip()
                if not message:
                    continue

                LOGGER.debug(f"← stdio: {message[:200]}...")

                # Notify handlers
                for handler in self._message_handlers:
                    try:
                        await handler(message)
                    except Exception as e:
                        LOGGER.error(f"Handler error: {e}")

        except asyncio.CancelledError:  # pylint: disable=try-except-raise
            raise
        except Exception as e:
            LOGGER.error(f"Error reading stdout: {e}")


class ReverseProxyClient:
    """Reverse proxy client that bridges local stdio to remote gateway."""

    def __init__(
        self,
        gateway_url: str,
        local_command: str,
        token: Optional[str] = None,
        reconnect_delay: float = DEFAULT_RECONNECT_DELAY,
        max_retries: int = DEFAULT_MAX_RETRIES,
        keepalive_interval: int = DEFAULT_KEEPALIVE_INTERVAL,
    ):
        """Initialize reverse proxy client.

        Args:
            gateway_url: Remote gateway URL.
            local_command: Local MCP server command.
            token: Optional bearer token for authentication.
            reconnect_delay: Initial reconnection delay in seconds.
            max_retries: Maximum reconnection attempts (0 = infinite).
            keepalive_interval: Heartbeat interval in seconds.
        """
        self.gateway_url = gateway_url
        self.local_command = local_command
        self.token = token
        self.reconnect_delay = reconnect_delay
        self.max_retries = max_retries
        self.keepalive_interval = keepalive_interval

        # Parse gateway URL
        parsed = urlparse(gateway_url)
        self.use_websocket = parsed.scheme in ("ws", "wss", "http", "https")

        # Connection state
        self.state = ConnectionState.DISCONNECTED
        self.connection: Optional[WSClientProtocol] = None
        self.session_id = uuid.uuid4().hex
        self.retry_count = 0

        # Components
        self.stdio_process = StdioProcess(local_command)
        self.stdio_process.add_message_handler(self._handle_stdio_message)

        # Tasks
        self._keepalive_task: Optional[asyncio.Task] = None
        self._receive_task: Optional[asyncio.Task] = None

        # Request tracking for correlation
        self._pending_requests: Dict[Any, asyncio.Future] = {}

    async def connect(self) -> None:
        """Establish connection to remote gateway.

        Raises:
            Exception: If connection fails.
        """
        if self.state != ConnectionState.DISCONNECTED:
            return

        self.state = ConnectionState.CONNECTING

        try:
            # Start local server first
            await self.stdio_process.start()

            # Connect to gateway
            if self.use_websocket:
                await self._connect_websocket()
            else:
                await self._connect_sse()

            self.state = ConnectionState.CONNECTED
            self.retry_count = 0

            # Register with gateway
            await self._register()

            # Start keepalive
            self._keepalive_task = asyncio.create_task(self._keepalive_loop())

            LOGGER.info(f"Connected to gateway: {self.gateway_url}")

        except Exception as e:
            LOGGER.error(f"Connection failed: {e}")
            self.state = ConnectionState.DISCONNECTED
            raise

    async def _connect_websocket(self) -> None:
        """Connect via WebSocket.

        Raises:
            ImportError: If websockets package is not installed.
        """
        if not websockets:
            raise ImportError("websockets package required for WebSocket support")

        # Build WebSocket URL
        ws_url = self.gateway_url.replace("http://", "ws://").replace("https://", "wss://")
        if not ws_url.startswith(("ws://", "wss://")):
            ws_url = f"wss://{ws_url}"

        # Add reverse proxy endpoint
        if "/reverse-proxy" not in ws_url:
            ws_url = urljoin(ws_url, "/reverse-proxy/ws")

        # Build headers
        headers = {}
        if self.token:
            headers["Authorization"] = f"Bearer {self.token}"
        headers["X-Session-ID"] = self.session_id

        LOGGER.info(f"Connecting to WebSocket: {ws_url}")

        # Connect
        self.connection = await websockets.connect(
            ws_url,
            extra_headers=headers,
            ping_interval=20,
            ping_timeout=10,
        )

        # Start receiving messages
        self._receive_task = asyncio.create_task(self._receive_websocket())

    async def _connect_sse(self) -> None:
        """Connect via SSE (fallback).

        Raises:
            ImportError: If httpx package is not installed.
            NotImplementedError: SSE transport not yet implemented.
        """
        if not httpx:
            raise ImportError("httpx package required for SSE support")

        # SSE implementation would establish SSE connection
        # and use POST endpoint for sending messages
        raise NotImplementedError("SSE transport not yet implemented")

    async def _register(self) -> None:
        """Register local server with gateway."""
        # Get server info by sending initialize request
        init_request = {
            "jsonrpc": "2.0",
            "id": "init-" + uuid.uuid4().hex,
            "method": "initialize",
            "params": {"protocolVersion": "2024-11-05", "capabilities": {}, "clientInfo": {"name": "reverse-proxy", "version": "1.0.0"}},
        }

        # Send to local server
        await self.stdio_process.send(json.dumps(init_request))

        # Wait for response (simplified - should correlate properly)
        await asyncio.sleep(1)

        # Send registration to gateway
        register_msg = {
            "type": MessageType.REGISTER.value,
            "sessionId": self.session_id,
            "server": {
                "name": f"reverse-proxy-{self.session_id[:8]}",
                "description": f"Reverse proxied: {self.local_command}",
                "protocol": "stdio",
            },
        }

        await self._send_to_gateway(json.dumps(register_msg))

    async def _send_to_gateway(self, message: str) -> None:
        """Send message to remote gateway.

        Args:
            message: Message to send.

        Raises:
            RuntimeError: If not connected to gateway.
            NotImplementedError: If SSE transport is used (not implemented).
        """
        conn = self.connection
        if not conn:
            raise RuntimeError("Not connected to gateway")

        if self.use_websocket:
            await cast(Any, conn).send(message)
        else:
            # SSE would POST to message endpoint
            raise NotImplementedError("SSE transport not yet implemented")

    async def _handle_stdio_message(self, message: str) -> None:
        """Handle message from local stdio server.

        Args:
            message: JSON-RPC message from stdio.
        """
        try:
            # Parse to check if it's a response or notification
            data = json.loads(message)

            # Wrap in reverse proxy envelope
            envelope = {"type": MessageType.RESPONSE.value if "id" in data else MessageType.NOTIFICATION.value, "sessionId": self.session_id, "payload": data}

            # Forward to gateway
            await self._send_to_gateway(json.dumps(envelope))

        except Exception as e:
            LOGGER.error(f"Error forwarding stdio message: {e}")

    async def _receive_websocket(self) -> None:
        """Receive messages from WebSocket connection."""
        if not self.connection:
            return

        try:
            conn = cast(Any, self.connection)
            async for message in conn:
                await self._handle_gateway_message(message)
        except Exception as e:  # Catch broad exceptions to avoid dependency-specific attribute errors
            closed_exc = None
            if websockets is not None:
                ex_mod = getattr(websockets, "exceptions", None)
                if ex_mod is not None:
                    closed_exc = getattr(ex_mod, "ConnectionClosed", None)
            if closed_exc and isinstance(e, closed_exc):
                LOGGER.warning("WebSocket connection closed")
            else:
                LOGGER.error(f"WebSocket receive error: {e}")
        finally:
            self.state = ConnectionState.DISCONNECTED

    async def _handle_gateway_message(self, message: str) -> None:
        """Handle message from remote gateway.

        Args:
            message: Message from gateway.
        """
        try:
            data = json.loads(message)
            msg_type = data.get("type")

            if msg_type == MessageType.REQUEST.value:
                # Forward request to local server
                payload = data.get("payload", {})
                await self.stdio_process.send(json.dumps(payload))

            elif msg_type == MessageType.HEARTBEAT.value:
                # Respond to heartbeat
                pong = {
                    "type": MessageType.HEARTBEAT.value,
                    "sessionId": self.session_id,
                }
                await self._send_to_gateway(json.dumps(pong))

            elif msg_type == MessageType.ERROR.value:
                LOGGER.error(f"Gateway error: {data.get('message', 'Unknown error')}")

            else:
                LOGGER.warning(f"Unknown message type: {msg_type}")

        except Exception as e:
            LOGGER.error(f"Error handling gateway message: {e}")

    async def _keepalive_loop(self) -> None:
        """Send periodic keepalive messages."""
        try:
            while self.state == ConnectionState.CONNECTED:
                await asyncio.sleep(self.keepalive_interval)

                heartbeat = {
                    "type": MessageType.HEARTBEAT.value,
                    "sessionId": self.session_id,
                }

                try:
                    await self._send_to_gateway(json.dumps(heartbeat))
                except Exception as e:
                    LOGGER.warning(f"Keepalive failed: {e}")
                    break

        except asyncio.CancelledError:
            pass

    async def disconnect(self) -> None:
        """Disconnect from gateway and stop local server."""
        if self.state == ConnectionState.SHUTTING_DOWN:
            return

        self.state = ConnectionState.SHUTTING_DOWN
        LOGGER.info("Disconnecting reverse proxy...")

        # Cancel tasks
        if self._keepalive_task:
            self._keepalive_task.cancel()
        if self._receive_task:
            self._receive_task.cancel()

        # Send unregister message
        if self.connection:
            try:
                unregister = {
                    "type": MessageType.UNREGISTER.value,
                    "sessionId": self.session_id,
                }
                await self._send_to_gateway(json.dumps(unregister))
            except Exception:
                pass  # nosec B110 - Intentionally swallow errors during cleanup

        # Close connection
        if self.connection:
            await cast(Any, self.connection).close()

        # Stop local server
        await self.stdio_process.stop()

        self.state = ConnectionState.DISCONNECTED
        LOGGER.info("Reverse proxy disconnected")

    async def run_with_reconnect(self) -> None:
        """Run the reverse proxy with automatic reconnection."""
        while True:
            try:
                if self.state == ConnectionState.SHUTTING_DOWN:
                    break

                await self.connect()

                # Wait for disconnection
                while self.state == ConnectionState.CONNECTED:
                    await asyncio.sleep(1)

                if self.state == ConnectionState.SHUTTING_DOWN:
                    break

            except Exception as e:
                LOGGER.error(f"Connection error: {e}")

            # Check retry limit
            self.retry_count += 1
            if self.max_retries > 0 and self.retry_count >= self.max_retries:
                LOGGER.error(f"Max retries ({self.max_retries}) exceeded")
                break

            # Calculate backoff delay
            delay = min(self.reconnect_delay * (2**self.retry_count), 60)
            LOGGER.info(f"Reconnecting in {delay}s (attempt {self.retry_count})")

            self.state = ConnectionState.RECONNECTING
            await asyncio.sleep(delay)
            self.state = ConnectionState.DISCONNECTED


def parse_args(argv: Optional[List[str]] = None) -> argparse.Namespace:
    """Parse command line arguments.

    Args:
        argv: Command line arguments (default: sys.argv[1:]).

    Returns:
        Parsed arguments.

    Examples:
        >>> import os
        >>> os.environ['REVERSE_PROXY_GATEWAY'] = 'https://example.com'
        >>> args = parse_args(['--local-stdio', 'mcp-server'])
        >>> args.local_stdio
        'mcp-server'
        >>> args.gateway
        'https://example.com'
        >>> args.log_level
        'INFO'
        >>> args = parse_args(['--local-stdio', 'mcp-server', '--verbose'])
        >>> args.log_level
        'DEBUG'
        >>> args = parse_args(['--local-stdio', 'mcp-server', '--max-retries', '5'])
        >>> args.max_retries
        5
    """
    parser = argparse.ArgumentParser(
        prog="mcpgateway.reverse_proxy",
        description="Bridge local MCP servers to remote gateways",
    )

    # Required arguments
    parser.add_argument(
        "--local-stdio",
        required=True,
        help="Local MCP server command to run via stdio",
    )

    parser.add_argument(
        "--gateway",
        help="Remote gateway URL (can also use REVERSE_PROXY_GATEWAY env var)",
    )

    # Authentication
    parser.add_argument(
        "--token",
        help="Bearer token for authentication (can also use REVERSE_PROXY_TOKEN env var)",
    )

    # Connection options
    parser.add_argument(
        "--reconnect-delay",
        type=float,
        default=DEFAULT_RECONNECT_DELAY,
        help=f"Initial reconnection delay in seconds (default: {DEFAULT_RECONNECT_DELAY})",
    )

    parser.add_argument(
        "--max-retries",
        type=int,
        default=DEFAULT_MAX_RETRIES,
        help=f"Maximum reconnection attempts, 0=infinite (default: {DEFAULT_MAX_RETRIES})",
    )

    parser.add_argument(
        "--keepalive",
        type=int,
        default=DEFAULT_KEEPALIVE_INTERVAL,
        help=f"Keepalive interval in seconds (default: {DEFAULT_KEEPALIVE_INTERVAL})",
    )

    # Configuration file
    parser.add_argument(
        "--config",
        help="Configuration file (YAML or JSON)",
    )

    # Logging
    parser.add_argument(
        "--log-level",
        choices=["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"],
        default="INFO",
        help="Log level (default: INFO)",
    )

    parser.add_argument(
        "--verbose",
        action="store_true",
        help="Enable verbose logging (same as --log-level DEBUG)",
    )

    args = parser.parse_args(argv)

    # Handle verbose flag
    if args.verbose:
        args.log_level = "DEBUG"

    # Get gateway from environment if not provided
    if not args.gateway:
        args.gateway = os.getenv(ENV_GATEWAY)
        if not args.gateway:
            parser.error("--gateway or REVERSE_PROXY_GATEWAY environment variable required")

    # Get token from environment if not provided
    if not args.token:
        args.token = os.getenv(ENV_TOKEN)

    # Load configuration file if provided
    if args.config:
        if not yaml:
            parser.error("PyYAML package required for configuration file support")
        yaml_module = cast(Any, yaml)

        with open(args.config, "r", encoding="utf-8") as f:
            if args.config.endswith((".yaml", ".yml")):
                config = yaml_module.safe_load(f)
            else:
                config = json.load(f)

        # Merge configuration (command line takes precedence)
        if not isinstance(config, dict):
            parser.error("Configuration file must contain a JSON/YAML object at the top level")
        else:
            for key, value in config.items():
                if not hasattr(args, key) or getattr(args, key) is None:
                    setattr(args, key, value)

    return args


async def main(argv: Optional[List[str]] = None) -> None:
    """Main entry point for reverse proxy.

    Args:
        argv: Command line arguments.
    """
    args = parse_args(argv)

    # Configure logging
    logging.basicConfig(
        level=getattr(logging, args.log_level),
        format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
        stream=sys.stderr,
    )

    # Create reverse proxy client
    client = ReverseProxyClient(
        gateway_url=args.gateway,
        local_command=args.local_stdio,
        token=args.token,
        reconnect_delay=args.reconnect_delay,
        max_retries=args.max_retries,
        keepalive_interval=args.keepalive,
    )

    # Handle shutdown signals
    shutdown_event = asyncio.Event()

    def signal_handler(*_args: object) -> None:
        """Handle shutdown signals gracefully.

        Args:
            *_args: Signal handler positional arguments (ignored).
        """
        LOGGER.info("Shutdown signal received")
        shutdown_event.set()

    # Register signal handlers
    loop = asyncio.get_running_loop()
    for sig in (signal.SIGINT, signal.SIGTERM):
        with suppress(NotImplementedError):
            loop.add_signal_handler(sig, signal_handler)

    # Run client with reconnection
    client_task = asyncio.create_task(client.run_with_reconnect())

    try:
        # Wait for shutdown
        await shutdown_event.wait()
    finally:
        # Clean shutdown
        await client.disconnect()
        client_task.cancel()
        with suppress(asyncio.CancelledError):
            await client_task


def run() -> None:
    """Console script entry point."""
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        print("\nShutdown complete", file=sys.stderr)
        sys.exit(0)
    except Exception as e:
        print(f"Error: {e}", file=sys.stderr)
        sys.exit(1)


if __name__ == "__main__":
    run()
