"""mTLS-enabled HTTP client wrapper for SecureMesh (TrustShield Pillar 4).

Provides an ``MTLSClient`` that wraps :class:`httpx.AsyncClient` with mutual
TLS authentication.  The client presents its own certificate (signed by the
local CA) and verifies that the server certificate is also signed by the same
CA and has not been revoked.

When mTLS is disabled (``enabled=False`` in config), the factory function
returns a plain :class:`httpx.AsyncClient` so callers don't need conditional
logic.
"""

from __future__ import annotations

import logging
import ssl
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any

import httpx

if TYPE_CHECKING:
    from llmhosts.security.tls_manager import TLSManager

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------


@dataclass(slots=True)
class MTLSConfig:
    """Configuration for an mTLS connection.

    Attributes
    ----------
    enabled:
        When *False*, mTLS is bypassed and a plain client is returned.
    tls_manager:
        The :class:`TLSManager` instance that owns the CA and certs.
    client_name:
        Certificate common name to use as the client identity
        (typically ``"proxy"``).
    verify_revocation:
        When *True*, check the revocation list before connecting.
    """

    enabled: bool = True
    tls_manager: TLSManager | None = None
    client_name: str = "proxy"
    verify_revocation: bool = True


# ---------------------------------------------------------------------------
# mTLS Client
# ---------------------------------------------------------------------------


class MTLSClient:
    """Async HTTP client with mutual TLS authentication.

    Wraps :class:`httpx.AsyncClient` so that every request presents the
    local client certificate and validates the server certificate against
    the local CA.

    Parameters
    ----------
    config:
        mTLS configuration including the :class:`TLSManager` and
        identity name.
    """

    def __init__(self, config: MTLSConfig) -> None:
        if config.tls_manager is None:
            raise ValueError("MTLSConfig.tls_manager must be set when mTLS is enabled")

        self._config = config
        self._tls = config.tls_manager
        self._client: httpx.AsyncClient | None = None

    # ------------------------------------------------------------------
    # Lifecycle
    # ------------------------------------------------------------------

    async def _ensure_client(self) -> httpx.AsyncClient:
        """Lazily build the httpx client with the TLS context."""
        if self._client is not None:
            return self._client

        ssl_ctx = self._build_ssl_context()
        self._client = httpx.AsyncClient(verify=ssl_ctx)
        return self._client

    def _build_ssl_context(self) -> ssl.SSLContext:
        """Build an ``ssl.SSLContext`` configured for mutual TLS."""
        ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
        ctx.minimum_version = ssl.TLSVersion.TLSv1_2

        # Load the local CA as the trust anchor
        ca_crt = self._tls.ca_cert_path()
        if not ca_crt.exists():
            raise FileNotFoundError("CA certificate not found. Run TLSManager.init_ca() first.")
        ctx.load_verify_locations(cafile=str(ca_crt))

        # Load the client certificate and key
        key_path, crt_path = self._tls.cert_paths(self._config.client_name)
        if not key_path.exists() or not crt_path.exists():
            raise FileNotFoundError(
                f"Client certificate for '{self._config.client_name}' not found. Run TLSManager.issue_cert() first."
            )
        ctx.load_cert_chain(certfile=str(crt_path), keyfile=str(key_path))

        return ctx

    # ------------------------------------------------------------------
    # Revocation check
    # ------------------------------------------------------------------

    def _check_revocation(self, cert_pem: bytes) -> None:
        """Raise if the given PEM certificate has been revoked.

        Parameters
        ----------
        cert_pem:
            PEM-encoded certificate bytes.

        Raises
        ------
        ConnectionRefusedError
            If the certificate's serial number is in the revocation list.
        """
        from cryptography import x509 as cx509

        cert = cx509.load_pem_x509_certificate(cert_pem)
        if self._tls.is_revoked(cert.serial_number):
            raise ConnectionRefusedError(f"Server certificate (serial={cert.serial_number}) has been revoked")

    # ------------------------------------------------------------------
    # HTTP methods
    # ------------------------------------------------------------------

    async def get(self, url: str, **kwargs: Any) -> httpx.Response:
        """Send an mTLS-authenticated GET request."""
        client = await self._ensure_client()
        return await client.get(url, **kwargs)

    async def post(self, url: str, **kwargs: Any) -> httpx.Response:
        """Send an mTLS-authenticated POST request."""
        client = await self._ensure_client()
        return await client.post(url, **kwargs)

    async def close(self) -> None:
        """Close the underlying HTTP client and release resources."""
        if self._client is not None:
            await self._client.aclose()
            self._client = None

    # ------------------------------------------------------------------
    # Context manager support
    # ------------------------------------------------------------------

    async def __aenter__(self) -> MTLSClient:
        await self._ensure_client()
        return self

    async def __aexit__(self, *exc: object) -> None:
        await self.close()


# ---------------------------------------------------------------------------
# Factory
# ---------------------------------------------------------------------------


def create_mtls_client(config: MTLSConfig) -> MTLSClient | httpx.AsyncClient:
    """Create an HTTP client with optional mTLS.

    When ``config.enabled`` is *False*, returns a plain
    :class:`httpx.AsyncClient` (no TLS client certs, standard verification).

    Parameters
    ----------
    config:
        mTLS configuration.

    Returns
    -------
    MTLSClient | httpx.AsyncClient
        An mTLS-wrapped client or a regular async HTTP client.
    """
    if not config.enabled:
        logger.info("mTLS disabled -- returning plain httpx.AsyncClient")
        return httpx.AsyncClient()

    logger.info("Creating mTLS client with identity '%s'", config.client_name)
    return MTLSClient(config)
