"""Minimal built-in SOCKS5 proxy server.

Runs inside MuraveyChina process — no external proxy tools needed.
Supports CONNECT command only (sufficient for HTTPS API calls).
Can forward through an upstream proxy or work as a direct SOCKS5 proxy.

Usage:
    from src.utils.socks5_server import MicroSocks5

    server = MicroSocks5(listen_port=11081)
    server.start()  # non-blocking, runs in daemon thread
    # proxy URL: socks5://127.0.0.1:11081
    ...
    server.stop()
"""

from __future__ import annotations

import logging
import select
import socket
import ssl
import struct
import threading

log = logging.getLogger(__name__)

# SOCKS5 constants
_VER = 0x05
_CMD_CONNECT = 0x01
_ATYP_IPV4 = 0x01
_ATYP_DOMAIN = 0x03
_ATYP_IPV6 = 0x04
_REP_SUCCESS = 0x00
_REP_GENERAL_FAILURE = 0x01
_REP_HOST_UNREACHABLE = 0x04
_REP_CMD_NOT_SUPPORTED = 0x07

_BUFFER_SIZE = 8192


class MicroSocks5:
    """Minimal SOCKS5 proxy server running in-process.

    Binds to localhost only (safe — not exposed to network).
    Each connection is handled in a separate daemon thread.
    """

    def __init__(
        self,
        listen_port: int = 11081,
        listen_host: str = "127.0.0.1",
        upstream_proxy: str = "",  # optional: forward through another proxy
        connect_timeout: float = 10.0,
    ) -> None:
        self.listen_host = listen_host
        self.listen_port = listen_port
        self.upstream_proxy = upstream_proxy
        self.connect_timeout = connect_timeout

        self._server_socket: socket.socket | None = None
        self._accept_thread: threading.Thread | None = None
        self._stop_event = threading.Event()
        self._active_connections: int = 0
        self._lock = threading.Lock()

    @property
    def proxy_url(self) -> str:
        return f"socks5://{self.listen_host}:{self.listen_port}"

    @property
    def is_running(self) -> bool:
        return self._server_socket is not None and not self._stop_event.is_set()

    @property
    def active_connections(self) -> int:
        with self._lock:
            return self._active_connections

    def start(self) -> bool:
        """Start the SOCKS5 server. Returns True on success."""
        if self.is_running:
            return True

        try:
            self._server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self._server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
            self._server_socket.settimeout(1.0)
            self._server_socket.bind((self.listen_host, self.listen_port))
            self._server_socket.listen(32)
        except OSError as exc:
            log.error("SOCKS5 server failed to bind %s:%d: %s", self.listen_host, self.listen_port, exc)
            self._server_socket = None
            return False

        self._stop_event.clear()
        self._accept_thread = threading.Thread(
            target=self._accept_loop,
            daemon=True,
            name="socks5-accept",
        )
        self._accept_thread.start()
        log.info("SOCKS5 server started on %s:%d", self.listen_host, self.listen_port)
        return True

    def stop(self) -> None:
        """Stop the SOCKS5 server."""
        self._stop_event.set()
        if self._server_socket:
            try:
                self._server_socket.close()
            except Exception:
                pass
            self._server_socket = None
        if self._accept_thread:
            self._accept_thread.join(timeout=3)
            self._accept_thread = None
        log.info("SOCKS5 server stopped")

    def _accept_loop(self) -> None:
        """Accept incoming connections."""
        while not self._stop_event.is_set():
            try:
                client_sock, addr = self._server_socket.accept()
                client_sock.settimeout(self.connect_timeout)
                t = threading.Thread(
                    target=self._handle_client,
                    args=(client_sock,),
                    daemon=True,
                    name=f"socks5-{addr[1]}",
                )
                t.start()
            except socket.timeout:
                continue
            except OSError:
                if not self._stop_event.is_set():
                    log.debug("SOCKS5 accept error")
                break

    def _handle_client(self, client: socket.socket) -> None:
        """Handle a single SOCKS5 client connection."""
        remote: socket.socket | None = None
        with self._lock:
            self._active_connections += 1
        try:
            # 1. Greeting: client sends version + auth methods
            data = client.recv(256)
            if len(data) < 3 or data[0] != _VER:
                return

            # Reply: no authentication required
            client.sendall(struct.pack("BB", _VER, 0x00))

            # 2. Connection request
            data = client.recv(256)
            if len(data) < 7 or data[0] != _VER:
                return

            cmd = data[1]
            if cmd != _CMD_CONNECT:
                self._send_reply(client, _REP_CMD_NOT_SUPPORTED)
                return

            atyp = data[3]

            if atyp == _ATYP_IPV4:
                dst_addr = socket.inet_ntoa(data[4:8])
                dst_port = struct.unpack("!H", data[8:10])[0]
            elif atyp == _ATYP_DOMAIN:
                domain_len = data[4]
                dst_addr = data[5:5 + domain_len].decode("ascii")
                dst_port = struct.unpack("!H", data[5 + domain_len:7 + domain_len])[0]
            elif atyp == _ATYP_IPV6:
                dst_addr = socket.inet_ntop(socket.AF_INET6, data[4:20])
                dst_port = struct.unpack("!H", data[20:22])[0]
            else:
                self._send_reply(client, _REP_GENERAL_FAILURE)
                return

            # 3. Connect to destination
            try:
                remote = socket.create_connection(
                    (dst_addr, dst_port), timeout=self.connect_timeout,
                )
                # Enable TCP keepalive on remote connection
                remote.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)
            except Exception:
                self._send_reply(client, _REP_HOST_UNREACHABLE)
                return

            # 4. Send success reply
            bind_addr = remote.getsockname()
            bind_ip = socket.inet_aton(bind_addr[0])
            bind_port = bind_addr[1]
            reply = struct.pack("!BBBB4sH", _VER, _REP_SUCCESS, 0x00, _ATYP_IPV4, bind_ip, bind_port)
            client.sendall(reply)

            # 5. Relay data bidirectionally
            self._relay(client, remote)

        except Exception as exc:
            log.debug("SOCKS5 handler error: %s", exc)
        finally:
            with self._lock:
                self._active_connections -= 1
            try:
                client.close()
            except Exception:
                pass
            if remote:
                try:
                    remote.close()
                except Exception:
                    pass

    def _send_reply(self, client: socket.socket, rep: int) -> None:
        """Send SOCKS5 reply."""
        reply = struct.pack("!BBBB4sH", _VER, rep, 0x00, _ATYP_IPV4, b"\x00\x00\x00\x00", 0)
        try:
            client.sendall(reply)
        except Exception:
            pass

    def _relay(self, client: socket.socket, remote: socket.socket) -> None:
        """Bidirectional data relay using select()."""
        client.setblocking(False)
        remote.setblocking(False)

        sockets = [client, remote]

        while not self._stop_event.is_set():
            try:
                readable, _, exceptional = select.select(sockets, [], sockets, 30.0)
            except (ValueError, OSError):
                break

            if exceptional:
                break

            if not readable:
                continue  # timeout — keep alive

            for sock in readable:
                try:
                    data = sock.recv(_BUFFER_SIZE)
                    if not data:
                        return  # Connection closed

                    other = remote if sock is client else client
                    other.sendall(data)
                except (BlockingIOError, ssl.SSLWantReadError):
                    continue
                except (OSError, BrokenPipeError):
                    return

    def __del__(self) -> None:
        self.stop()
