"""In-process network guard for replay mode.

Blocks ALL network access during replay to ensure determinism:
- socket.create_connection, socket.connect, socket.connect_ex
- socket.getaddrinfo (DNS resolution)
- socket.socket.sendto (UDP bypass)
- ssl.SSLContext.wrap_socket (TLS connections)
- asyncio.open_connection (async TCP)
"""

from __future__ import annotations

import socket
from typing import Any

from autodebug_replay.transcripts import ReplayDivergence


class NetworkGuard:
    def __init__(self) -> None:
        self._orig_create_connection = socket.create_connection
        self._orig_getaddrinfo = socket.getaddrinfo
        self._orig_connect = socket.socket.connect
        self._orig_connect_ex = socket.socket.connect_ex
        self._orig_sendto = socket.socket.sendto
        self._orig_ssl_wrap = None
        self._orig_ssl_ctx_wrap = None
        self._orig_asyncio_open = None

    def _fail(self, detail: dict) -> None:
        raise ReplayDivergence("NETWORK_ACCESS", detail)

    def create_connection(self, *args: Any, **kwargs: Any):
        self._fail({"op": "create_connection"})

    def getaddrinfo(self, *args: Any, **kwargs: Any):
        self._fail({"op": "getaddrinfo"})

    def connect(self, *args: Any, **kwargs: Any):
        self._fail({"op": "connect"})

    def connect_ex(self, *args: Any, **kwargs: Any):
        self._fail({"op": "connect_ex"})

    def sendto(self, *args: Any, **kwargs: Any):
        self._fail({"op": "sendto"})

    def ssl_context_wrap(self, *args: Any, **kwargs: Any):
        self._fail({"op": "ssl_context_wrap_socket"})

    def install(self) -> None:
        socket.create_connection = self.create_connection  # type: ignore[assignment]
        socket.getaddrinfo = self.getaddrinfo  # type: ignore[assignment]
        socket.socket.connect = self.connect  # type: ignore[assignment]
        socket.socket.connect_ex = self.connect_ex  # type: ignore[assignment]
        socket.socket.sendto = self.sendto  # type: ignore[assignment]

        # Block SSL
        try:
            import ssl
            self._orig_ssl_ctx_wrap = ssl.SSLContext.wrap_socket
            ssl.SSLContext.wrap_socket = self.ssl_context_wrap  # type: ignore[assignment]
            if hasattr(ssl, "wrap_socket"):
                self._orig_ssl_wrap = ssl.wrap_socket
                ssl.wrap_socket = self.ssl_context_wrap  # type: ignore[assignment]
        except ImportError:
            pass

        # Block asyncio transports
        try:
            import asyncio

            orig_open = asyncio.open_connection
            self._orig_asyncio_open = orig_open
            guard = self

            async def blocked_open_connection(*a, **kw):
                guard._fail({"op": "asyncio.open_connection"})

            asyncio.open_connection = blocked_open_connection  # type: ignore[assignment]
        except ImportError:
            pass

    def uninstall(self) -> None:
        """Restore all original network functions."""
        socket.create_connection = self._orig_create_connection  # type: ignore[assignment]
        socket.getaddrinfo = self._orig_getaddrinfo  # type: ignore[assignment]
        socket.socket.connect = self._orig_connect  # type: ignore[assignment]
        socket.socket.connect_ex = self._orig_connect_ex  # type: ignore[assignment]
        socket.socket.sendto = self._orig_sendto  # type: ignore[assignment]

        try:
            import ssl
            if self._orig_ssl_ctx_wrap is not None:
                ssl.SSLContext.wrap_socket = self._orig_ssl_ctx_wrap  # type: ignore[assignment]
            if self._orig_ssl_wrap is not None and hasattr(ssl, "wrap_socket"):
                ssl.wrap_socket = self._orig_ssl_wrap  # type: ignore[assignment]
        except ImportError:
            pass

        try:
            import asyncio
            if self._orig_asyncio_open is not None:
                asyncio.open_connection = self._orig_asyncio_open  # type: ignore[assignment]
        except ImportError:
            pass
