########################################################################################################################
# IMPORTS

import contextlib
import logging
import select
import socket
import struct
import threading

import socks

########################################################################################################################
# SETUP

logger = logging.getLogger(__name__)

########################################################################################################################
# CLASS


class Socks5Bridge:
    """
    Local SOCKS5 proxy (no auth) that forwards TCP through an authenticated SOCKS5 proxy.

    Playwright's Firefox does not support SOCKS5 proxy authentication, but it does support
    unauthenticated SOCKS5. This bridge forwards TCP CONNECT requests to the upstream
    SOCKS5 proxy with credentials via PySocks.

    Firefox sees ``socks5://127.0.0.1:{port}`` with no credentials required.

    Usage:
        bridge = Socks5Bridge(host, port, username, password)
        local_port = bridge.start()
        # use socks5://127.0.0.1:{local_port} as proxy for Playwright
        bridge.stop()
    """

    # SOCKS5
    _VER = 0x05
    _NO_AUTH = 0x00
    _CMD_CONNECT = 0x01
    _ATYP_IPV4 = 0x01
    _ATYP_DOMAIN = 0x03
    _ATYP_IPV6 = 0x04
    _REP_SUCCESS = 0x00
    _REP_FAILURE = 0x01
    _CMD_NOT_SUPPORTED = 0x07

    def __init__(self, host: str, port: int, username: str | None = None, password: str | None = None):
        self.host = host
        self.port = int(port)
        self.username = username
        self.password = password
        self._server: socket.socket | None = None
        self._local_port: int | None = None
        self._running = False
        self._thread: threading.Thread | None = None

    # ------------------------------------------------------------------------------------------------------------------
    # Public API

    def start(self) -> int:
        """Start the bridge, return the local port."""
        self._server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self._server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        self._server.bind(("127.0.0.1", 0))
        self._server.settimeout(1)
        self._local_port = self._server.getsockname()[1]
        self._server.listen(50)
        self._running = True
        self._thread = threading.Thread(target=self._accept_loop, daemon=True, name="socks5-bridge")
        self._thread.start()
        logger.info(f"SOCKS5 bridge listening on 127.0.0.1:{self._local_port} → {self.host}:{self.port}")
        return self._local_port

    def stop(self) -> None:
        self._running = False
        if self._server:
            with contextlib.suppress(Exception):
                self._server.close()
        if self._thread:
            self._thread.join(timeout=2)

    # ------------------------------------------------------------------------------------------------------------------
    # TCP accept loop

    def _accept_loop(self) -> None:
        while self._running:
            try:
                conn, _ = self._server.accept()
                threading.Thread(target=self._handle, args=(conn,), daemon=True).start()
            except socket.timeout:
                continue
            except Exception:
                break

    # ------------------------------------------------------------------------------------------------------------------
    # SOCKS5 handshake dispatcher

    def _handle(self, conn: socket.socket) -> None:
        try:
            header = self._recvall(conn, 2)
            if len(header) < 2 or header[0] != self._VER:
                logger.warning(f"SOCKS5 bridge: bad greeting {header!r}")
                return
            self._recvall(conn, header[1])  # discard offered methods
            conn.sendall(bytes([self._VER, self._NO_AUTH]))

            req = self._recvall(conn, 4)
            if len(req) < 4 or req[0] != self._VER:
                logger.warning(f"SOCKS5 bridge: bad request header {req!r}")
                conn.sendall(self._reply(self._REP_FAILURE))
                return

            cmd, atyp = req[1], req[3]
            if cmd == self._CMD_CONNECT:
                self._handle_connect(conn, atyp)
            else:
                logger.warning(f"SOCKS5 bridge: unsupported cmd={cmd}")
                conn.sendall(self._reply(self._CMD_NOT_SUPPORTED))
        except Exception as ex:
            logger.warning(f"SOCKS5 bridge handle error: {ex}")
        finally:
            with contextlib.suppress(Exception):
                conn.close()

    # ------------------------------------------------------------------------------------------------------------------
    # TCP CONNECT → upstream SOCKS5

    def _handle_connect(self, conn: socket.socket, atyp: int) -> None:
        dest_host, dest_port = self._read_addr_port(conn, atyp)
        logger.debug(f"SOCKS5 bridge CONNECT → {dest_host}:{dest_port}")
        remote = None
        try:
            remote = socks.socksocket()
            remote.set_proxy(socks.SOCKS5, self.host, self.port, rdns=True,
                             username=self.username, password=self.password)
            remote.settimeout(30)
            remote.connect((dest_host, dest_port))
            conn.sendall(self._reply(self._REP_SUCCESS))
            self._relay(conn, remote)
        except Exception as ex:
            logger.warning(f"SOCKS5 bridge CONNECT error ({dest_host}:{dest_port}): {ex}")
            with contextlib.suppress(Exception):
                conn.sendall(self._reply(self._REP_FAILURE))
        finally:
            if remote:
                with contextlib.suppress(Exception):
                    remote.close()

    # ------------------------------------------------------------------------------------------------------------------
    # Helpers

    @staticmethod
    def _recvall(conn: socket.socket, n: int) -> bytes:
        buf = b""
        while len(buf) < n:
            chunk = conn.recv(n - len(buf))
            if not chunk:
                raise ConnectionError("Connection closed before receiving expected bytes")
            buf += chunk
        return buf

    def _read_addr_port(self, conn: socket.socket, atyp: int) -> tuple[str, int]:
        if atyp == self._ATYP_IPV4:
            host = socket.inet_ntoa(self._recvall(conn, 4))
        elif atyp == self._ATYP_DOMAIN:
            length = self._recvall(conn, 1)[0]
            host = self._recvall(conn, length).decode()
        elif atyp == self._ATYP_IPV6:
            host = socket.inet_ntop(socket.AF_INET6, self._recvall(conn, 16))
        else:
            raise ValueError(f"Unknown ATYP: {atyp}")
        port = struct.unpack("!H", self._recvall(conn, 2))[0]
        return host, port

    @staticmethod
    def _reply(rep: int) -> bytes:
        return struct.pack("!BBBBIH", 0x05, rep, 0x00, 0x01, 0, 0)

    @staticmethod
    def _relay(client: socket.socket, remote: socket.socket) -> None:
        while True:
            try:
                r, _, _ = select.select([client, remote], [], [], 30)
            except Exception:
                break
            if not r:
                break
            for s in r:
                try:
                    data = s.recv(8192)
                    if not data:
                        return
                    (remote if s is client else client).sendall(data)
                except Exception:
                    return
