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

Unix socket client for external plugins.

This module provides a high-performance client for communicating with
external plugins over Unix domain sockets using length-prefixed protobuf
messages.

Examples:
    Create and use a Unix socket plugin client:

    >>> from mcpgateway.plugins.framework.external.unix.client import UnixSocketExternalPlugin
    >>> from mcpgateway.plugins.framework.models import PluginConfig, UnixSocketClientConfig

    >>> config = PluginConfig(
    ...     name="MyPlugin",
    ...     kind="external",
    ...     hooks=["tool_pre_invoke"],
    ...     unix_socket=UnixSocketClientConfig(path="/tmp/plugin.sock"),
    ... )
    >>> plugin = UnixSocketExternalPlugin(config)
    >>> # await plugin.initialize()
    >>> # result = await plugin.invoke_hook(hook_type, payload, context)
"""
# pylint: disable=no-member,no-name-in-module

# Standard
import asyncio
import logging
from typing import Any, Optional

# Third-Party
from google.protobuf import json_format
from google.protobuf.struct_pb2 import Struct

# First-Party
from mcpgateway.plugins.framework.base import Plugin
from mcpgateway.plugins.framework.errors import convert_exception_to_error, PluginError
from mcpgateway.plugins.framework.external.grpc.proto import plugin_service_pb2
from mcpgateway.plugins.framework.external.proto_convert import pydantic_context_to_proto, update_pydantic_context_from_proto
from mcpgateway.plugins.framework.external.unix.protocol import read_message, write_message_async
from mcpgateway.plugins.framework.hooks.registry import get_hook_registry
from mcpgateway.plugins.framework.models import PluginConfig, PluginContext, PluginErrorModel, PluginResult

logger = logging.getLogger(__name__)


class UnixSocketExternalPlugin(Plugin):
    """External plugin client using raw Unix domain sockets.

    This client provides high-performance IPC for local plugins using
    length-prefixed protobuf messages. It includes automatic reconnection
    with configurable retry logic.

    Attributes:
        config: The plugin configuration.

    Examples:
        >>> from mcpgateway.plugins.framework.models import PluginConfig, UnixSocketClientConfig
        >>> config = PluginConfig(
        ...     name="TestPlugin",
        ...     kind="external",
        ...     hooks=["tool_pre_invoke"],
        ...     unix_socket=UnixSocketClientConfig(path="/tmp/test.sock"),
        ... )
        >>> plugin = UnixSocketExternalPlugin(config)
        >>> plugin.name
        'TestPlugin'
    """

    def __init__(self, config: PluginConfig) -> None:
        """Initialize the Unix socket plugin client.

        Args:
            config: The plugin configuration with unix_socket settings.

        Raises:
            PluginError: If unix_socket configuration is missing.
        """
        super().__init__(config)

        if not config.unix_socket:
            raise PluginError(error=PluginErrorModel(message="The unix_socket section must be defined for Unix socket plugin", plugin_name=config.name))

        self._socket_path = config.unix_socket.path
        self._reconnect_attempts = config.unix_socket.reconnect_attempts
        self._reconnect_delay = config.unix_socket.reconnect_delay
        self._timeout = config.unix_socket.timeout

        self._reader: Optional[asyncio.StreamReader] = None
        self._writer: Optional[asyncio.StreamWriter] = None
        self._connected = False
        self._lock = asyncio.Lock()

    @property
    def connected(self) -> bool:
        """Check if the client is connected.

        Returns:
            bool: True if connected and writer is active, False otherwise.
        """
        return self._connected and self._writer is not None and not self._writer.is_closing()

    async def _connect(self) -> None:
        """Establish connection to the Unix socket server.

        Raises:
            PluginError: If connection fails.
        """
        try:
            self._reader, self._writer = await asyncio.open_unix_connection(self._socket_path)
            self._connected = True
            logger.debug("Connected to Unix socket: %s", self._socket_path)
        except OSError as e:
            self._connected = False
            raise PluginError(error=PluginErrorModel(message=f"Failed to connect to {self._socket_path}: {e}", plugin_name=self.name)) from e

    async def _disconnect(self) -> None:
        """Close the connection."""
        if self._writer:
            try:
                self._writer.close()
                await self._writer.wait_closed()
            except Exception:  # nosec B110 - cleanup code, exceptions should not propagate
                pass
        self._writer = None
        self._reader = None
        self._connected = False

    async def _reconnect(self) -> None:
        """Attempt to reconnect with retry logic.

        Raises:
            PluginError: If all reconnection attempts fail.
        """
        await self._disconnect()

        last_error: Optional[Exception] = None
        for attempt in range(1, self._reconnect_attempts + 1):
            try:
                logger.debug("Reconnection attempt %d/%d to %s", attempt, self._reconnect_attempts, self._socket_path)
                await self._connect()
                logger.info("Reconnected to %s on attempt %d", self._socket_path, attempt)
                return
            except PluginError as e:
                last_error = e
                if attempt < self._reconnect_attempts:
                    await asyncio.sleep(self._reconnect_delay * attempt)  # Exponential backoff

        raise PluginError(error=PluginErrorModel(message=f"Failed to reconnect after {self._reconnect_attempts} attempts: {last_error}", plugin_name=self.name))

    async def _send_request(self, request: plugin_service_pb2.InvokeHookRequest) -> plugin_service_pb2.InvokeHookResponse:
        """Send a request and receive response, with reconnection on failure.

        Args:
            request: The protobuf request to send.

        Returns:
            The protobuf response.

        Raises:
            PluginError: If sending fails after reconnection attempts.
        """
        request_bytes = request.SerializeToString()

        async with self._lock:
            for attempt in range(self._reconnect_attempts + 1):
                try:
                    if not self.connected:
                        await self._reconnect()

                    # Send request
                    await write_message_async(self._writer, request_bytes)

                    # Read response
                    response_bytes = await read_message(self._reader, timeout=self._timeout)

                    # Parse response
                    response = plugin_service_pb2.InvokeHookResponse()
                    response.ParseFromString(response_bytes)
                    return response

                except asyncio.TimeoutError as e:
                    logger.warning("Request timed out after %s seconds", self._timeout)
                    raise PluginError(error=PluginErrorModel(message=f"Request timed out after {self._timeout}s", plugin_name=self.name)) from e

                except (OSError, asyncio.IncompleteReadError, BrokenPipeError) as e:
                    logger.warning("Connection error on attempt %d: %s", attempt + 1, e)
                    self._connected = False

                    if attempt < self._reconnect_attempts:
                        await asyncio.sleep(self._reconnect_delay * (attempt + 1))
                        continue
                    raise PluginError(error=PluginErrorModel(message=f"Request failed after {self._reconnect_attempts + 1} attempts: {e}", plugin_name=self.name)) from e

        # Should not reach here
        raise PluginError(error=PluginErrorModel(message="Unexpected state in _send_request", plugin_name=self.name))

    async def initialize(self) -> None:
        """Initialize the plugin client by connecting to the server.

        This establishes the Unix socket connection and optionally
        fetches the remote plugin configuration.

        Raises:
            PluginError: If initial connection fails.
        """
        logger.info("Initializing Unix socket plugin: %s -> %s", self.name, self._socket_path)

        try:
            await self._connect()
        except PluginError:
            raise
        except Exception as e:
            logger.exception(e)
            raise PluginError(error=convert_exception_to_error(e, plugin_name=self.name))

        # Optionally fetch remote config to verify connection
        try:
            request = plugin_service_pb2.GetPluginConfigRequest(name=self.name)
            request_bytes = request.SerializeToString()

            await write_message_async(self._writer, request_bytes)
            response_bytes = await read_message(self._reader, timeout=self._timeout)

            response = plugin_service_pb2.GetPluginConfigResponse()
            response.ParseFromString(response_bytes)

            if response.found:
                logger.debug("Remote plugin config verified for %s", self.name)
            else:
                logger.warning("Plugin %s not found on remote server", self.name)

        except Exception as e:
            logger.warning("Could not verify remote plugin config: %s", e)
            # Continue anyway - the plugin might still work

        logger.info("Unix socket plugin initialized: %s", self.name)

    async def shutdown(self) -> None:
        """Shutdown the plugin client and close the connection."""
        logger.info("Shutting down Unix socket plugin: %s", self.name)
        await self._disconnect()

    async def invoke_hook(
        self,
        hook_type: str,
        payload: Any,
        context: PluginContext,
    ) -> PluginResult:
        """Invoke a plugin hook over the Unix socket connection.

        Args:
            hook_type: The type of hook to invoke (e.g., "tool_pre_invoke").
            payload: The hook payload (will be serialized to protobuf Struct).
            context: The plugin context.

        Returns:
            The plugin result.

        Raises:
            PluginError: If the request fails after retries or hook type is invalid.
        """
        # 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))

        # Convert payload to Struct (still polymorphic)
        payload_struct = Struct()
        if hasattr(payload, "model_dump"):
            json_format.ParseDict(payload.model_dump(), payload_struct)
        else:
            json_format.ParseDict(payload, payload_struct)

        # Convert context to explicit proto message (faster than Struct)
        context_proto = pydantic_context_to_proto(context)

        # Build request
        request = plugin_service_pb2.InvokeHookRequest(
            hook_type=hook_type,
            plugin_name=self.name,
            payload=payload_struct,
            context=context_proto,
        )

        try:
            # Send request and get response
            response = await self._send_request(request)

            # Handle error response
            if response.HasField("error") and response.error.message:
                error = PluginErrorModel(
                    message=response.error.message,
                    plugin_name=response.error.plugin_name or self.name,
                    code=response.error.code,
                    mcp_error_code=response.error.mcp_error_code,
                )
                if response.error.HasField("details"):
                    error.details = json_format.MessageToDict(response.error.details)
                raise PluginError(error=error)

            # Update context if modified (using explicit proto message)
            if response.HasField("context"):
                update_pydantic_context_from_proto(context, response.context)

            # Parse and return result
            if response.HasField("result"):
                result_dict = json_format.MessageToDict(response.result)
                return result_type.model_validate(result_dict)

            raise PluginError(
                error=PluginErrorModel(
                    message="Received invalid response from Unix socket plugin server",
                    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))
