from __future__ import annotations

import inspect
import socket
import threading
from functools import lru_cache
from importlib import import_module
from typing import Any

from govoplan_core.security.outbound_http import (
    OutboundHttpError,
    create_outbound_connection,
)


class SdkPeerPinningError(RuntimeError):
    """Raised when an optional SDK cannot be bound to the pinned transport."""


_BOTOCORE_CLIENT_CREATION_LOCK = threading.Lock()


def create_pinned_s3_client(**kwargs: Any) -> Any:
    """Construct an S3 client whose first and every later socket is pinned.

    Botocore does not expose the HTTP-session class through boto3's public
    client API. Its endpoint creator does expose that seam, so this function
    replaces the creator only for the bounded client-construction operation.
    The lock avoids an unsafe interleaving with another GovOPlaN client
    construction. Any unrelated botocore client created during the short
    replacement window also receives the stricter transport.
    """

    try:
        boto3 = import_module("boto3")
        botocore_args = import_module("botocore.args")
    except ImportError as exc:
        raise SdkPeerPinningError(
            "S3 connector browsing requires the optional boto3 dependency"
        ) from exc

    pinned_session_cls = _botocore_transport_types()[0]
    original_creator = getattr(botocore_args, "EndpointCreator", None)
    if original_creator is None or not hasattr(original_creator, "create_endpoint"):
        raise SdkPeerPinningError(
            "The installed botocore release does not expose the endpoint transport seam required for peer pinning"
        )

    class PinnedEndpointCreator(original_creator):  # type: ignore[misc, valid-type]
        def create_endpoint(self, *args: Any, **endpoint_kwargs: Any) -> Any:
            endpoint_kwargs["http_session_cls"] = pinned_session_cls
            return super().create_endpoint(*args, **endpoint_kwargs)

    with _BOTOCORE_CLIENT_CREATION_LOCK:
        current_creator = getattr(botocore_args, "EndpointCreator", None)
        if current_creator is not original_creator:
            raise SdkPeerPinningError(
                "Botocore endpoint construction changed concurrently; refusing to create an unproven S3 client"
            )
        botocore_args.EndpointCreator = PinnedEndpointCreator
        try:
            session_type = getattr(getattr(boto3, "session", None), "Session", None)
            if session_type is None:
                raise SdkPeerPinningError(
                    "The installed boto3 release does not expose its isolated session constructor"
                )
            client = session_type().client("s3", **kwargs)
        finally:
            if getattr(botocore_args, "EndpointCreator", None) is PinnedEndpointCreator:
                botocore_args.EndpointCreator = original_creator

    http_session = getattr(getattr(client, "_endpoint", None), "http_session", None)
    if http_session is None or not isinstance(http_session, pinned_session_cls):
        close = getattr(client, "close", None)
        if callable(close):
            close()
        raise SdkPeerPinningError(
            "Botocore did not install the required pinned HTTP transport; the S3 client was discarded"
        )
    return client


@lru_cache(maxsize=1)
def _botocore_transport_types() -> tuple[type[Any], type[Any], type[Any]]:
    try:
        awsrequest = import_module("botocore.awsrequest")
        httpsession = import_module("botocore.httpsession")
        urllib3_exceptions = import_module("urllib3.exceptions")
    except ImportError as exc:
        raise SdkPeerPinningError(
            "S3 connector browsing requires a compatible botocore HTTP transport"
        ) from exc

    required = (
        "AWSHTTPConnection",
        "AWSHTTPSConnection",
        "AWSHTTPConnectionPool",
        "AWSHTTPSConnectionPool",
    )
    if any(not hasattr(awsrequest, name) for name in required) or not hasattr(
        httpsession, "URLLib3Session"
    ):
        raise SdkPeerPinningError(
            "The installed botocore release is missing the connection classes required for peer pinning"
        )

    def pinned_new_connection(connection: Any) -> socket.socket:
        hostname = str(getattr(connection, "_dns_host", "") or "").strip()
        port = int(getattr(connection, "port", 0) or 0)
        if not hostname or not port:
            raise urllib3_exceptions.NewConnectionError(
                connection, "Pinned S3 connection is missing its target authority"
            )
        try:
            return create_outbound_connection(
                hostname,
                port,
                timeout=getattr(connection, "timeout", None),
                source_address=getattr(connection, "source_address", None),
                socket_options=getattr(connection, "socket_options", None),
                label="S3 connector peer",
            )
        except socket.timeout as exc:
            raise urllib3_exceptions.ConnectTimeoutError(
                connection,
                f"Connection to {hostname} timed out while selecting an approved peer",
            ) from exc
        except (OSError, OutboundHttpError, ValueError) as exc:
            raise urllib3_exceptions.NewConnectionError(
                connection,
                f"S3 connector peer was rejected: {exc}",
            ) from exc

    class PinnedAWSHTTPConnection(awsrequest.AWSHTTPConnection):
        _new_conn = pinned_new_connection

    class PinnedAWSHTTPSConnection(awsrequest.AWSHTTPSConnection):
        _new_conn = pinned_new_connection

    class PinnedAWSHTTPConnectionPool(awsrequest.AWSHTTPConnectionPool):
        ConnectionCls = PinnedAWSHTTPConnection

    class PinnedAWSHTTPSConnectionPool(awsrequest.AWSHTTPSConnectionPool):
        ConnectionCls = PinnedAWSHTTPSConnection

    class PinnedURLLib3Session(httpsession.URLLib3Session):
        def __init__(self, *args: Any, **kwargs: Any) -> None:
            # A proxy would select the final target outside this process. Until
            # proxy peer delegation is modeled, connector traffic is direct.
            kwargs["proxies"] = {}
            super().__init__(*args, **kwargs)
            self._pool_classes_by_scheme = {
                "http": PinnedAWSHTTPConnectionPool,
                "https": PinnedAWSHTTPSConnectionPool,
            }
            manager = getattr(self, "_manager", None)
            if manager is None or not hasattr(manager, "pool_classes_by_scheme"):
                raise SdkPeerPinningError(
                    "The installed botocore pool manager cannot enforce pinned connection classes"
                )
            manager.pool_classes_by_scheme = self._pool_classes_by_scheme

    return PinnedURLLib3Session, PinnedAWSHTTPConnection, PinnedAWSHTTPSConnection


@lru_cache(maxsize=1)
def pinned_smb_connection_cache() -> dict[str, Any]:
    """Return the Files-owned cache; unproven process-global sessions are never reused."""

    return {}


def install_pinned_smb_transport(smbclient: Any) -> Any:
    """Install a process-wide fail-closed smbclient session factory.

    smbclient routes initial connections, reconnects, and DFS targets through
    ``smbclient._pool.register_session``. Replacing that single factory with a
    behavior-compatible implementation ensures every target uses a socket
    selected by Core's deployment-wide outbound policy.
    """

    try:
        pool = import_module("smbclient._pool")
        connection_module = import_module("smbprotocol.connection")
        session_module = import_module("smbprotocol.session")
        transport_module = import_module("smbprotocol.transport")
    except ImportError as exc:
        raise SdkPeerPinningError(
            "SMB connector browsing requires the optional smbprotocol dependency"
        ) from exc

    current = getattr(pool, "register_session", None)
    if getattr(current, "__govoplan_peer_pinned__", False):
        expected_tcp = getattr(current, "__govoplan_pinned_tcp__", None)
        if expected_tcp is None or getattr(connection_module, "Tcp", None) is not expected_tcp:
            raise SdkPeerPinningError(
                "The installed SMB transport changed after peer pinning; refusing to reuse the session factory"
            )
        return smbclient
    required_parameters = {
        "server",
        "username",
        "password",
        "port",
        "encrypt",
        "connection_timeout",
        "connection_cache",
        "auth_protocol",
        "require_signing",
    }
    if current is None or not required_parameters.issubset(
        inspect.signature(current).parameters
    ):
        raise SdkPeerPinningError(
            "The installed smbprotocol release does not expose the session seam required for peer pinning"
        )

    class PinnedTcp(transport_module.Tcp):
        def connect(self) -> None:
            with self._sock_lock:
                if self.connected:
                    return
                try:
                    self._sock = create_outbound_connection(
                        self.server,
                        int(self.port),
                        timeout=self.timeout,
                        label="SMB connector peer",
                    )
                except (OSError, OutboundHttpError, ValueError) as exc:
                    raise ValueError(
                        f"SMB connector peer '{self.server}:{self.port}' was rejected: {exc}"
                    ) from exc
                self._sock.settimeout(None)
                self.connected = True

    # Connection.connect() instantiates the module-level Tcp symbol on every
    # reconnect. Replacing it process-wide is deliberate: an SMB connection
    # created by another code path must become stricter, never bypass Files'
    # peer boundary.
    if not hasattr(connection_module, "Tcp"):
        raise SdkPeerPinningError(
            "The installed smbprotocol release cannot install the pinned TCP transport"
        )
    connection_module.Tcp = PinnedTcp

    def pinned_register_session(
        server: str,
        username: str | None = None,
        password: str | None = None,
        port: int = 445,
        encrypt: bool | None = None,
        connection_timeout: float = 60,
        connection_cache: dict[str, Any] | None = None,
        auth_protocol: str = "negotiate",
        require_signing: bool = True,
    ) -> Any:
        cache = pinned_smb_connection_cache() if connection_cache is None else connection_cache
        connection_key = f"{server.lower()}:{port}"
        connection = cache.get(connection_key)
        transport = getattr(connection, "transport", None)
        if connection is not None and not isinstance(transport, PinnedTcp):
            disconnect = getattr(connection, "disconnect", None)
            if callable(disconnect):
                try:
                    disconnect(close=True)
                except Exception:
                    pass
            cache.pop(connection_key, None)
            connection = None
        if connection is None or not getattr(connection.transport, "connected", False):
            connection = connection_module.Connection(
                pool.ClientConfig().client_guid,
                server,
                port,
                require_signing=require_signing,
            )
            connection.transport = PinnedTcp(server, port)
            connection.connect(timeout=connection_timeout)
            if not isinstance(connection.transport, PinnedTcp):
                disconnect = getattr(connection, "disconnect", None)
                if callable(disconnect):
                    try:
                        disconnect(close=True)
                    except Exception:
                        pass
                raise SdkPeerPinningError(
                    "smbprotocol replaced the required pinned TCP transport during connection setup"
                )
            cache[connection_key] = connection

        session = next(
            (
                item
                for item in connection.session_table.values()
                if username is None or item.username == username
            ),
            None,
        )
        if session is None:
            session = session_module.Session(
                connection,
                username=username,
                password=password,
                require_encryption=(encrypt is True),
                auth_protocol=auth_protocol,
            )
            session.connect()
        elif encrypt is not None:
            if session.encrypt_data and not encrypt:
                raise ValueError(
                    "Cannot disable encryption on an already negotiated session."
                )
            if not session.encrypt_data and encrypt:
                session.encrypt = True
        return session

    pinned_register_session.__govoplan_peer_pinned__ = True  # type: ignore[attr-defined]
    pinned_register_session.__govoplan_pinned_tcp__ = PinnedTcp  # type: ignore[attr-defined]
    pool.register_session = pinned_register_session
    if hasattr(smbclient, "register_session"):
        smbclient.register_session = pinned_register_session
    return smbclient


__all__ = [
    "SdkPeerPinningError",
    "create_pinned_s3_client",
    "install_pinned_smb_transport",
    "pinned_smb_connection_cache",
]