# -*- coding: utf-8 -*-
"""Location: ./mcpgateway/plugins/framework/external/mcp/client.py
Copyright 2025
SPDX-License-Identifier: Apache-2.0
Authors: Teryl Taylor

External plugin client which connects to a remote server through MCP.
Module that contains plugin MCP client code to serve external plugins.
"""

# Standard
import asyncio
from contextlib import AsyncExitStack
from functools import partial
import logging
import os
from pathlib import Path
import sys
from typing import Any, Awaitable, Callable, Optional

# Third-Party
import httpx
from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client
from mcp.client.streamable_http import streamablehttp_client
from mcp.types import TextContent
import orjson

# First-Party
from mcpgateway.common.models import TransportType
from mcpgateway.config import settings
from mcpgateway.plugins.framework.base import HookRef, Plugin, PluginRef
from mcpgateway.plugins.framework.constants import CONTEXT, ERROR, GET_PLUGIN_CONFIG, HOOK_TYPE, IGNORE_CONFIG_EXTERNAL, INVOKE_HOOK, NAME, PAYLOAD, PLUGIN_NAME, PYTHON_SUFFIX, RESULT
from mcpgateway.plugins.framework.errors import convert_exception_to_error, PluginError
from mcpgateway.plugins.framework.external.mcp.tls_utils import create_ssl_context
from mcpgateway.plugins.framework.hooks.registry import get_hook_registry
from mcpgateway.plugins.framework.models import MCPClientTLSConfig, PluginConfig, PluginContext, PluginErrorModel, PluginPayload, PluginResult

logger = logging.getLogger(__name__)


class ExternalPlugin(Plugin):
    """External plugin object for pre/post processing of inputs and outputs at various locations throughout the gateway.

    The External Plugin connects to a remote MCP server that contains plugins.
    """

    def __init__(self, config: PluginConfig) -> None:
        """Initialize a plugin with a configuration and context.

        Args:
            config: The plugin configuration
        """
        super().__init__(config)
        self._session: Optional[ClientSession] = None
        self._exit_stack = AsyncExitStack()
        self._http: Optional[Any]
        self._stdio: Optional[Any]
        self._write: Optional[Any]
        self._current_task = asyncio.current_task()
        self._stdio_exit_stack: Optional[AsyncExitStack] = None
        self._stdio_task: Optional[asyncio.Task[None]] = None
        self._stdio_ready: Optional[asyncio.Event] = None
        self._stdio_stop: Optional[asyncio.Event] = None
        self._stdio_error: Optional[BaseException] = None
        self._get_session_id: Optional[Callable[[], str | None]] = None
        self._session_id: Optional[str] = None
        self._http_client_factory: Optional[Callable[..., httpx.AsyncClient]] = None

    async def initialize(self) -> None:
        """Initialize the plugin's connection to the MCP server.

        Raises:
            PluginError: if unable to retrieve plugin configuration of external plugin.
        """

        if not self._config.mcp:
            raise PluginError(error=PluginErrorModel(message="The mcp section must be defined for external plugin", plugin_name=self.name))
        if self._config.mcp.proto == TransportType.STDIO:
            if not (self._config.mcp.script or self._config.mcp.cmd):
                raise PluginError(error=PluginErrorModel(message="STDIO transport requires script or cmd", plugin_name=self.name))
            await self.__connect_to_stdio_server(self._config.mcp.script, self._config.mcp.cmd, self._config.mcp.env, self._config.mcp.cwd)
        elif self._config.mcp.proto == TransportType.STREAMABLEHTTP:
            if not self._config.mcp.url:
                raise PluginError(error=PluginErrorModel(message="STREAMABLEHTTP transport requires url", plugin_name=self.name))
            await self.__connect_to_http_server(self._config.mcp.url)

        try:
            config = await self.__get_plugin_config()

            if not config:
                raise PluginError(error=PluginErrorModel(message="Unable to retrieve configuration for external plugin", plugin_name=self.name))

            current_config = self._config.model_dump(exclude_unset=True)
            remote_config = config.model_dump(exclude_unset=True)
            remote_config.update(current_config)

            context = {IGNORE_CONFIG_EXTERNAL: True}

            self._config = PluginConfig.model_validate(remote_config, context=context)
        except PluginError as pe:
            try:
                await self.shutdown()
            except Exception as shutdown_error:
                logger.error("Error during external plugin shutdown after init failure: %s", shutdown_error)
            logger.exception(pe)
            raise
        except Exception as e:
            try:
                await self.shutdown()
            except Exception as shutdown_error:
                logger.error("Error during external plugin shutdown after init failure: %s", shutdown_error)
            logger.exception(e)
            raise PluginError(error=convert_exception_to_error(e, plugin_name=self.name))

    def __resolve_stdio_command(self, script_path: str | None, cmd: list[str] | None, cwd: str | None) -> tuple[str, list[str]]:
        """Resolve the stdio command + args from config.

        Args:
            script_path: Path to a server script or executable.
            cmd: Command list to execute (command + args).
            cwd: Working directory for resolving relative script paths.

        Returns:
            Tuple of (command, args).

        Raises:
            PluginError: if the script is invalid or cmd is malformed.
        """
        if cmd:
            if not isinstance(cmd, list) or not cmd or not all(isinstance(part, str) and part.strip() for part in cmd):
                raise PluginError(error=PluginErrorModel(message="STDIO cmd must be a non-empty list of strings", plugin_name=self.name))
            return cmd[0], cmd[1:]

        if not script_path:
            raise PluginError(error=PluginErrorModel(message="STDIO transport requires script or cmd", plugin_name=self.name))

        server_path = Path(script_path).expanduser()
        if not server_path.is_absolute() and cwd:
            server_path = Path(cwd).expanduser() / server_path
        resolved_script_path = str(server_path)
        if not server_path.is_file():
            raise PluginError(error=PluginErrorModel(message=f"Server script {resolved_script_path} does not exist.", plugin_name=self.name))

        if server_path.suffix == PYTHON_SUFFIX:
            return sys.executable, [resolved_script_path]
        if server_path.suffix == ".sh":
            return "sh", [resolved_script_path]
        if not os.access(server_path, os.X_OK):
            raise PluginError(error=PluginErrorModel(message=f"Server script {resolved_script_path} must be executable.", plugin_name=self.name))
        return resolved_script_path, []

    def __build_stdio_env(self, extra_env: dict[str, str] | None) -> dict[str, str]:
        """Build environment for the stdio server process.

        Args:
            extra_env: Environment overrides to merge into the current process env.

        Returns:
            Combined environment dictionary for the plugin process.
        """
        current_env = os.environ.copy()
        if extra_env:
            current_env.update(extra_env)
        return current_env

    async def __run_stdio_session(self, server_script_path: str | None, cmd: list[str] | None, env: dict[str, str] | None, cwd: str | None) -> None:
        """Run a stdio session in a dedicated task for consistent setup/teardown.

        Args:
            server_script_path: Path to the server script or executable.
            cmd: Command list to start the server (command + args).
            env: Environment overrides for the server process.
            cwd: Working directory for the server process.
        """
        try:
            command, args = self.__resolve_stdio_command(server_script_path, cmd, cwd)
            server_env = self.__build_stdio_env(env)
            server_params = StdioServerParameters(command=command, args=args, env=server_env, cwd=cwd)

            self._stdio_exit_stack = AsyncExitStack()
            stdio_transport = await self._stdio_exit_stack.enter_async_context(stdio_client(server_params))
            self._stdio, self._write = stdio_transport
            self._session = await self._stdio_exit_stack.enter_async_context(ClientSession(self._stdio, self._write))

            await self._session.initialize()

            response = await self._session.list_tools()
            tools = response.tools
            logger.info("\nConnected to plugin MCP server (stdio) with tools: %s", " ".join([tool.name for tool in tools]))
        except Exception as e:
            self._stdio_error = e
            logger.exception(e)
        finally:
            if self._stdio_ready and not self._stdio_ready.is_set():
                self._stdio_ready.set()

        if self._stdio_error:
            if self._stdio_exit_stack:
                await self._stdio_exit_stack.aclose()
            return

        if self._stdio_stop:
            await self._stdio_stop.wait()

        if self._stdio_exit_stack:
            await self._stdio_exit_stack.aclose()

    async def __connect_to_stdio_server(self, server_script_path: str | None, cmd: list[str] | None, env: dict[str, str] | None, cwd: str | None) -> None:
        """Connect to an MCP plugin server via stdio.

        Args:
            server_script_path: Path to the server script or executable.
            cmd: Command list to start the server (command + args).
            env: Environment overrides for the server process.
            cwd: Working directory for the server process.

        Raises:
            PluginError: if stdio script/cmd is invalid or if there is a connection error.
        """
        try:
            if not self._stdio_ready:
                self._stdio_ready = asyncio.Event()
            if not self._stdio_stop:
                self._stdio_stop = asyncio.Event()
            self._stdio_error = None

            self._stdio_task = asyncio.create_task(
                self.__run_stdio_session(server_script_path, cmd, env, cwd),
                name=f"external-plugin-stdio-{self.name}",
            )

            await self._stdio_ready.wait()
            if self._stdio_error:
                raise PluginError(error=convert_exception_to_error(self._stdio_error, plugin_name=self.name))
        except PluginError:
            raise
        except Exception as e:
            logger.exception(e)
            raise PluginError(error=convert_exception_to_error(e, plugin_name=self.name))

    async def __connect_to_http_server(self, uri: str) -> None:
        """Connect to an MCP plugin server via streamable http with retry logic.

        Args:
            uri: the URI of the mcp plugin server.

        Raises:
            PluginError: if there is an external connection error after all retries.
        """
        plugin_tls = self._config.mcp.tls if self._config and self._config.mcp else None
        uds_path = self._config.mcp.uds if self._config and self._config.mcp else None
        if uds_path and plugin_tls:
            logger.warning("TLS configuration is ignored for Unix domain socket connections.")
        tls_config = None if uds_path else (plugin_tls or MCPClientTLSConfig.from_env())

        def _tls_httpx_client_factory(
            headers: Optional[dict[str, str]] = None,
            timeout: Optional[httpx.Timeout] = None,
            auth: Optional[httpx.Auth] = None,
        ) -> httpx.AsyncClient:
            """Build an httpx client with TLS configuration for external MCP servers.

            Args:
                headers: Optional HTTP headers to include in requests.
                timeout: Optional timeout configuration for HTTP requests.
                auth: Optional authentication handler for HTTP requests.

            Returns:
                Configured httpx AsyncClient with TLS settings applied.

            Raises:
                PluginError: If TLS configuration fails.
            """

            # First-Party
            from mcpgateway.services.http_client_service import get_default_verify, get_http_timeout  # pylint: disable=import-outside-toplevel

            kwargs: dict[str, Any] = {"follow_redirects": True}
            if uds_path:
                kwargs["transport"] = httpx.AsyncHTTPTransport(uds=uds_path)
            if headers:
                kwargs["headers"] = headers
            kwargs["timeout"] = timeout if timeout else get_http_timeout()
            if auth is not None:
                kwargs["auth"] = auth

            # Add connection pool limits
            kwargs["limits"] = httpx.Limits(
                max_connections=settings.httpx_max_connections,
                max_keepalive_connections=settings.httpx_max_keepalive_connections,
                keepalive_expiry=settings.httpx_keepalive_expiry,
            )

            if not tls_config:
                # Use skip_ssl_verify setting when no custom TLS config
                kwargs["verify"] = get_default_verify()
                return httpx.AsyncClient(**kwargs)

            # Create SSL context using the utility function
            # This implements certificate validation per test_client_certificate_validation.py
            ssl_context = create_ssl_context(tls_config, self.name)
            kwargs["verify"] = ssl_context

            return httpx.AsyncClient(**kwargs)

        self._http_client_factory = _tls_httpx_client_factory
        max_retries = 3
        base_delay = 1.0

        for attempt in range(max_retries):

            try:
                client_factory = _tls_httpx_client_factory
                streamable_client = streamablehttp_client(uri, httpx_client_factory=client_factory, terminate_on_close=False)
                http_transport = await self._exit_stack.enter_async_context(streamable_client)
                self._http, self._write, get_session_id = http_transport
                self._get_session_id = get_session_id
                self._session = await self._exit_stack.enter_async_context(ClientSession(self._http, self._write))

                await self._session.initialize()
                self._session_id = self._get_session_id() if self._get_session_id else None
                response = await self._session.list_tools()
                tools = response.tools
                logger.info(
                    "Successfully connected to plugin MCP server with tools: %s",
                    " ".join([tool.name for tool in tools]),
                )
                return
            except Exception as e:
                logger.warning(f"Connection attempt {attempt + 1}/{max_retries} failed: {e}")
                if attempt == max_retries - 1:
                    # Final attempt failed
                    target = f"{uri} (uds={uds_path})" if uds_path else uri
                    error_msg = f"External plugin '{self.name}' connection failed after {max_retries} attempts: {target} is not reachable. Please ensure the MCP server is running."
                    logger.error(error_msg)
                    raise PluginError(error=PluginErrorModel(message=error_msg, plugin_name=self.name))
                await self.shutdown()
                # Wait before retry
                delay = base_delay * (2**attempt)
                logger.info(f"Retrying in {delay}s...")
                await asyncio.sleep(delay)

    async def invoke_hook(self, hook_type: str, payload: PluginPayload, context: PluginContext) -> PluginResult:
        """Invoke an external plugin hook using the MCP protocol.

        Args:
            hook_type:  The type of hook invoked (i.e., prompt_pre_fetch)
            payload: The payload to be passed to the hook.
            context: The plugin context passed to the run.

        Raises:
            PluginError: error passed from external plugin server.

        Returns:
            The resulting payload from the plugin.
        """
        # Get the result type from the global registry
        registry = get_hook_registry()
        result_type = registry.get_result_type(hook_type)
        if not result_type:
            raise PluginError(error=PluginErrorModel(message=f"Hook type '{hook_type}' not registered in hook registry", plugin_name=self.name))

        if not self._session:
            raise PluginError(error=PluginErrorModel(message="Plugin session not initialized", plugin_name=self.name))

        try:
            result = await self._session.call_tool(INVOKE_HOOK, {HOOK_TYPE: hook_type, PLUGIN_NAME: self.name, PAYLOAD: payload, CONTEXT: context})
            for content in result.content:
                if not isinstance(content, TextContent):
                    continue
                try:
                    res = orjson.loads(content.text)
                except orjson.JSONDecodeError:
                    raise PluginError(error=PluginErrorModel(message=f"Error trying to decode json: {content.text}", code="JSON_DECODE_ERROR", plugin_name=self.name))
                if CONTEXT in res:
                    cxt = PluginContext.model_validate(res[CONTEXT])
                    context.state = cxt.state
                    context.metadata = cxt.metadata
                    context.global_context.state = cxt.global_context.state
                if RESULT in res:
                    return result_type.model_validate(res[RESULT])
                if ERROR in res:
                    error = PluginErrorModel.model_validate(res[ERROR])
                    raise PluginError(error)
        except PluginError as pe:
            logger.exception(pe)
            raise
        except Exception as e:
            logger.exception(e)
            raise PluginError(error=convert_exception_to_error(e, plugin_name=self.name))
        raise PluginError(error=PluginErrorModel(message=f"Received invalid response. Result = {result}", plugin_name=self.name))

    async def __get_plugin_config(self) -> PluginConfig | None:
        """Retrieve plugin configuration for the current plugin on the remote MCP server.

        Raises:
            PluginError: if there is a connection issue or validation issue.

        Returns:
            A plugin configuration for the current plugin from a remote MCP server.
        """
        if not self._session:
            raise PluginError(error=PluginErrorModel(message="Plugin session not initialized", plugin_name=self.name))
        try:
            configs = await self._session.call_tool(GET_PLUGIN_CONFIG, {NAME: self.name})
            for content in configs.content:
                if not isinstance(content, TextContent):
                    continue
                conf = orjson.loads(content.text)
                if not conf:
                    return None
                return PluginConfig.model_validate(conf)
        except Exception as e:
            logger.exception(e)
            raise PluginError(error=convert_exception_to_error(e, plugin_name=self.name))

        return None

    async def shutdown(self) -> None:
        """Plugin cleanup code."""
        if self._stdio_task:
            if self._stdio_stop:
                self._stdio_stop.set()
            try:
                await self._stdio_task
            except Exception as e:
                logger.error("Error shutting down stdio session for plugin %s: %s", self.name, e)
            self._stdio_task = None
            self._stdio_ready = None
            self._stdio_stop = None
            self._stdio_exit_stack = None
            self._stdio_error = None
            self._stdio = None
            self._write = None
            if self._config and self._config.mcp and self._config.mcp.proto == TransportType.STDIO:
                self._session = None

        if self._exit_stack:
            await self._exit_stack.aclose()
        if self._config and self._config.mcp and self._config.mcp.proto == TransportType.STREAMABLEHTTP:
            await self.__terminate_http_session()
        self._get_session_id = None
        self._session_id = None
        self._http_client_factory = None

    async def __terminate_http_session(self) -> None:
        """Terminate streamable HTTP session explicitly to avoid lingering server state."""
        if not self._session_id or not self._config or not self._config.mcp or not self._config.mcp.url:
            return
        # Third-Party
        from mcp.server.streamable_http import MCP_SESSION_ID_HEADER  # pylint: disable=import-outside-toplevel

        client_factory = self._http_client_factory
        try:
            if client_factory:
                client = client_factory()
            else:
                client = httpx.AsyncClient(follow_redirects=True)
            async with client:
                headers = {MCP_SESSION_ID_HEADER: self._session_id}
                await client.delete(self._config.mcp.url, headers=headers)
        except Exception as exc:
            logger.debug("Failed to terminate streamable HTTP session: %s", exc)


class ExternalHookRef(HookRef):
    """A Hook reference point for external plugins."""

    def __init__(self, hook: str, plugin_ref: PluginRef):  # pylint: disable=super-init-not-called
        """Initialize a hook reference point for an external plugin.

        Note: We intentionally don't call super().__init__() because external plugins
        use invoke_hook() rather than direct method attributes.

        Args:
            hook: name of the hook point.
            plugin_ref: The reference to the plugin to hook.

        Raises:
            PluginError: If the plugin is not an external plugin.
        """
        self._plugin_ref = plugin_ref
        self._hook = hook
        if hasattr(plugin_ref.plugin, INVOKE_HOOK):
            self._func: Callable[[PluginPayload, PluginContext], Awaitable[PluginResult]] = partial(plugin_ref.plugin.invoke_hook, hook)  # type: ignore[attr-defined]
        else:
            raise PluginError(error=PluginErrorModel(message=f"Plugin: {plugin_ref.plugin.name} is not an external plugin", plugin_name=plugin_ref.plugin.name))
