"""Small in-memory HTTP/SOCKS proxy used by one browser generation."""

from __future__ import annotations

import base64
import select
import socket
import socketserver
import ssl
import threading
from collections.abc import Callable
from contextlib import suppress
from dataclasses import dataclass
from typing import cast
from urllib.parse import urlsplit


@dataclass(frozen=True)
class ProxyExit:
    protocol: str
    host: str
    port: int
    username: str = ""
    credential: str = ""


class _ThreadingTCPServer(socketserver.ThreadingMixIn, socketserver.TCPServer):
    allow_reuse_address = True
    daemon_threads = True


class ProxyRegistry:
    def __init__(self) -> None:
        self._lock = threading.RLock()
        self._proxies: dict[str, MemoryProxy] = {}

    def configure(
        self,
        alias: str,
        binding_version: int,
        bind_host: str,
        port: int,
        exit: ProxyExit,
        network_id: str = "",
    ) -> tuple[str, Callable[[], None]]:
        with self._lock:
            current = self._proxies.get(alias)
            if current and current.matches(
                binding_version, bind_host, port, exit, network_id
            ):
                return current.url, lambda: self._remove_object(alias, current)
            if current:
                self._proxies.pop(alias, None)
                current.close()
            proxy = MemoryProxy(
                alias, binding_version, bind_host, port, exit, network_id
            )
            self._proxies[alias] = proxy
        return proxy.url, lambda: self._remove_object(alias, proxy)

    def bind(
        self,
        alias: str,
        binding_version: int,
        proxy_url: str,
        runtime_id: str,
        network_id: str = "",
    ) -> bool:
        with self._lock:
            proxy = self._proxies.get(alias)
            return bool(
                proxy and proxy.bind(binding_version, proxy_url, runtime_id, network_id)
            )

    def ready(
        self, alias: str, port: int, runtime_id: str, network_id: str = ""
    ) -> bool:
        with self._lock:
            proxy = self._proxies.get(alias)
            return bool(proxy and proxy.ready(port, runtime_id, network_id))

    def remove(
        self, alias: str, binding_version: int, runtime_id: str, network_id: str = ""
    ) -> bool:
        with self._lock:
            proxy = self._proxies.get(alias)
            if not proxy:
                return True
            if not proxy.identity_matches(binding_version, runtime_id, network_id):
                return False
            self._proxies.pop(alias, None)
        proxy.close()
        return True

    def _remove_object(self, alias: str, proxy: MemoryProxy) -> None:
        with self._lock:
            if self._proxies.get(alias) is not proxy:
                return
            self._proxies.pop(alias, None)
        proxy.close()

    def close(self) -> None:
        with self._lock:
            proxies = list(self._proxies.values())
            self._proxies.clear()
        for proxy in proxies:
            proxy.close()


class MemoryProxy:
    def __init__(
        self,
        alias: str,
        binding_version: int,
        bind_host: str,
        port: int,
        exit: ProxyExit,
        network_id: str,
    ) -> None:
        self.alias = alias
        self.binding_version = binding_version
        self.bind_host = bind_host
        self.exit = exit
        self.network_id = network_id
        self.runtime_id = ""
        self._lock = threading.RLock()
        self._tunnels: set[socket.socket] = set()
        handler_type = type(
            "ProxyHandler",
            (_ProxyHandler,),
            {"proxy": self},
        )
        self.server = _ThreadingTCPServer((bind_host, port), handler_type)
        self.listener = self.server.socket
        actual_port = self.listener.getsockname()[1]
        self.url = f"http://{bind_host}:{actual_port}"
        self._thread = threading.Thread(
            target=self.server.serve_forever,
            name=f"creatorhub-proxy-{alias}",
            daemon=True,
        )
        self._thread.start()

    def matches(
        self,
        binding_version: int,
        bind_host: str,
        port: int,
        exit: ProxyExit,
        network_id: str,
    ) -> bool:
        actual_port = self.listener.getsockname()[1]
        return (
            self.binding_version == binding_version
            and self.bind_host == bind_host
            and (port == 0 or port == actual_port)
            and self.exit == exit
            and self.network_id == network_id
        )

    def bind(
        self, binding_version: int, proxy_url: str, runtime_id: str, network_id: str
    ) -> bool:
        with self._lock:
            if (
                binding_version != self.binding_version
                or proxy_url != self.url
                or not runtime_id
            ):
                return False
            self.runtime_id = runtime_id
            if network_id:
                self.network_id = network_id
            return True

    def ready(self, port: int, runtime_id: str, network_id: str) -> bool:
        with self._lock:
            return bool(
                runtime_id
                and self.runtime_id == runtime_id
                and self.listener.getsockname()[1] == port
                and (not network_id or self.network_id == network_id)
            )

    def identity_matches(
        self, binding_version: int, runtime_id: str, network_id: str
    ) -> bool:
        with self._lock:
            return (
                self.binding_version == binding_version
                and bool(runtime_id)
                and self.runtime_id == runtime_id
                and (not network_id or self.network_id == network_id)
            )

    def close(self) -> None:
        self.server.shutdown()
        self.server.server_close()
        with self._lock:
            tunnels = list(self._tunnels)
            self._tunnels.clear()
        for connection in tunnels:
            with suppress(OSError):
                connection.close()

    def add_tunnel(self, connection: socket.socket) -> bool:
        with self._lock:
            if self.server.socket.fileno() < 0:
                return False
            self._tunnels.add(connection)
            return True

    def remove_tunnel(self, connection: socket.socket) -> None:
        with self._lock:
            self._tunnels.discard(connection)

    def dial(self, target: str, timeout: float = 20.0) -> socket.socket:
        exit = self.exit
        if exit.protocol in ("http", "https"):
            return _dial_http_proxy(exit, target, timeout)
        if exit.protocol == "socks4":
            return _dial_socks4(exit, target, timeout)
        if exit.protocol == "socks5":
            return _dial_socks5(exit, target, timeout)
        raise OSError("unsupported proxy protocol")

    def forward_http(
        self,
        client: socket.socket,
        method: str,
        target: str,
        headers: list[tuple[str, str]],
        body: bytes,
    ) -> None:
        parsed = urlsplit(target)
        if not parsed.hostname:
            raise OSError("proxy request target is invalid")
        target_host = parsed.hostname
        target_port = parsed.port or (443 if parsed.scheme == "https" else 80)
        origin_target = parsed.path or "/"
        if parsed.query:
            origin_target += "?" + parsed.query
        if self.exit.protocol in ("http", "https"):
            upstream = _open_host(
                self.exit.host, self.exit.port, 20.0, self.exit.protocol == "https"
            )
            request_target = target
            request_headers = _forward_headers(headers, len(body))
            if self.exit.username:
                token = base64.b64encode(
                    f"{self.exit.username}:{self.exit.credential}".encode()
                ).decode()
                request_headers.append(("Proxy-Authorization", f"Basic {token}"))
        else:
            upstream = self.dial(f"{target_host}:{target_port}")
            request_target = origin_target
            request_headers = _forward_headers(headers, len(body))
        try:
            lines = [f"{method} {request_target} HTTP/1.1"]
            lines.extend(f"{k}: {v}" for k, v in request_headers)
            lines.append("Connection: close")
            upstream.sendall(("\r\n".join(lines) + "\r\n\r\n").encode() + body)
            _copy_until_close(upstream, client)
        finally:
            upstream.close()

    def tunnel(self, client: socket.socket, target: str) -> None:
        upstream = self.dial(target)
        try:
            client.sendall(b"HTTP/1.1 200 Connection Established\r\n\r\n")
            if not self.add_tunnel(client):
                return
            _relay(client, upstream)
        finally:
            self.remove_tunnel(client)
            with suppress(OSError):
                upstream.close()


class _ProxyHandler(socketserver.BaseRequestHandler):
    proxy: MemoryProxy

    def handle(self) -> None:
        client = self.request
        client.settimeout(60.0)
        try:
            head, body = _read_request(client)
            method, target, headers = _parse_request(head)
            if method.upper() == "CONNECT":
                self.proxy.tunnel(client, target)
            else:
                self.proxy.forward_http(client, method, target, headers, body)
        except (OSError, ValueError):
            with suppress(OSError):
                client.sendall(
                    b"HTTP/1.1 502 Bad Gateway\r\nConnection: close\r\nContent-Length: 0\r\n\r\n"
                )


def _read_request(connection: socket.socket) -> tuple[bytes, bytes]:
    data = bytearray()
    while b"\r\n\r\n" not in data:
        chunk = connection.recv(65536)
        if not chunk:
            raise OSError("proxy request closed")
        data.extend(chunk)
        if len(data) > 256 * 1024:
            raise OSError("proxy headers too large")
    split = data.index(b"\r\n\r\n") + 4
    head, initial_body = bytes(data[:split]), bytearray(data[split:])
    fields = _header_fields(head)
    lengths = fields.get("content-length", [])
    transfer_encoding = fields.get("transfer-encoding", [])
    if lengths and transfer_encoding:
        raise OSError("proxy request has both content length and transfer encoding")
    if transfer_encoding:
        if (
            len(transfer_encoding) != 1
            or transfer_encoding[0].lower().strip() != "chunked"
        ):
            raise OSError("unsupported proxy transfer encoding")
        return head, _read_chunked_body(connection, initial_body)
    content_length = 0
    if lengths:
        if len(lengths) != 1:
            raise OSError("proxy content length is ambiguous")
        try:
            content_length = int(lengths[0])
        except ValueError as exc:
            raise OSError("invalid proxy content length") from exc
    if content_length < 0 or content_length > 16 * 1024 * 1024:
        raise OSError("proxy body too large")
    while len(initial_body) < content_length:
        chunk = connection.recv(min(65536, content_length - len(initial_body)))
        if not chunk:
            raise OSError("proxy body closed")
        initial_body.extend(chunk)
    return head, bytes(initial_body[:content_length])


def _header_fields(head: bytes) -> dict[str, list[str]]:
    fields: dict[str, list[str]] = {}
    for line in head.decode("iso-8859-1").split("\r\n")[1:]:
        if not line:
            continue
        if ":" not in line:
            raise OSError("invalid proxy header")
        key, value = line.split(":", 1)
        fields.setdefault(key.strip().lower(), []).append(value.strip())
    return fields


def _read_line(connection: socket.socket, buffer: bytearray, limit: int) -> bytes:
    while b"\r\n" not in buffer:
        chunk = connection.recv(65536)
        if not chunk:
            raise OSError("proxy chunked body closed")
        buffer.extend(chunk)
        if len(buffer) > limit:
            raise OSError("proxy chunk line too large")
    end = buffer.index(b"\r\n")
    line = bytes(buffer[:end])
    del buffer[: end + 2]
    return line


def _read_chunked_body(connection: socket.socket, buffer: bytearray) -> bytes:
    body = bytearray()
    while True:
        line = _read_line(connection, buffer, 8192)
        size_text = line.split(b";", 1)[0].strip()
        try:
            size = int(size_text, 16)
        except ValueError as exc:
            raise OSError("invalid proxy chunk size") from exc
        if size < 0:
            raise OSError("invalid proxy chunk size")
        if size == 0:
            while True:
                trailer = _read_line(connection, buffer, 256 * 1024)
                if not trailer:
                    return bytes(body)
                if b":" not in trailer:
                    raise OSError("invalid proxy trailer")
        if len(body) + size > 16 * 1024 * 1024:
            raise OSError("proxy body too large")
        while len(buffer) < size + 2:
            chunk = connection.recv(65536)
            if not chunk:
                raise OSError("proxy chunked body closed")
            buffer.extend(chunk)
        body.extend(buffer[:size])
        if buffer[size : size + 2] != b"\r\n":
            raise OSError("invalid proxy chunk terminator")
        del buffer[: size + 2]


def _forward_headers(
    headers: list[tuple[str, str]], body_length: int
) -> list[tuple[str, str]]:
    return [
        (k, v)
        for k, v in headers
        if k.lower()
        not in (
            "proxy-authorization",
            "connection",
            "transfer-encoding",
            "content-length",
        )
    ] + [("Content-Length", str(body_length))]


def _parse_request(head: bytes) -> tuple[str, str, list[tuple[str, str]]]:
    lines = head.decode("iso-8859-1").split("\r\n")
    method, target, version = lines[0].split(" ", 2)
    if version not in ("HTTP/1.0", "HTTP/1.1") or not target:
        raise ValueError("invalid proxy request")
    headers: list[tuple[str, str]] = []
    for line in lines[1:]:
        if not line:
            continue
        if ":" not in line:
            raise ValueError("invalid proxy header")
        key, value = line.split(":", 1)
        headers.append((key.strip(), value.strip()))
    if method.upper() == "CONNECT":
        if ":" not in target:
            raise ValueError("CONNECT target missing port")
        return method, target, headers
    return method, target, headers


def _open_host(
    host: str, port: int, timeout: float, tls: bool = False
) -> socket.socket:
    connection = socket.create_connection((host, port), timeout=timeout)
    if tls:
        context = ssl.create_default_context()
        connection = context.wrap_socket(connection, server_hostname=host)
    return connection


def _dial_http_proxy(exit: ProxyExit, target: str, timeout: float) -> socket.socket:
    connection = _open_host(exit.host, exit.port, timeout, exit.protocol == "https")
    try:
        lines = [f"CONNECT {target} HTTP/1.1", f"Host: {target}"]
        if exit.username:
            token = base64.b64encode(
                f"{exit.username}:{exit.credential}".encode()
            ).decode()
            lines.append(f"Proxy-Authorization: Basic {token}")
        connection.sendall(("\r\n".join(lines) + "\r\n\r\n").encode())
        status, buffered = _read_connect_response(connection)
        if status != 200:
            raise OSError(f"upstream proxy returned {status}")
        result = cast(
            socket.socket,
            _BufferedSocket(connection, buffered) if buffered else connection,
        )
        result.settimeout(None)
        return result
    except Exception:
        connection.close()
        raise


def _dial_socks4(exit: ProxyExit, target: str, timeout: float) -> socket.socket:
    try:
        host, port_text = target.rsplit(":", 1)
        port = int(port_text)
    except (ValueError, IndexError) as exc:
        raise OSError("invalid SOCKS4 target") from exc
    if not 1 <= port <= 65535:
        raise OSError("invalid SOCKS4 target port")
    connection = _open_host(exit.host, exit.port, timeout)
    try:
        ip = socket.inet_aton(host) if _is_ipv4(host) else b"\x00\x00\x00\x01"
        payload = (
            b"\x04\x01"
            + port.to_bytes(2, "big")
            + ip
            + exit.username.encode()
            + b"\x00"
        )
        if ip == b"\x00\x00\x00\x01":
            payload += host.encode() + b"\x00"
        connection.sendall(payload)
        response = _recv_exact(connection, 8)
        if response[1] != 90:
            raise OSError("SOCKS4 proxy rejected connection")
        connection.settimeout(None)
        return connection
    except Exception:
        connection.close()
        raise


def _dial_socks5(exit: ProxyExit, target: str, timeout: float) -> socket.socket:
    try:
        host, port_text = target.rsplit(":", 1)
        port = int(port_text)
    except (ValueError, IndexError) as exc:
        raise OSError("invalid SOCKS5 target") from exc
    if not 1 <= port <= 65535:
        raise OSError("invalid SOCKS5 target port")
    connection = _open_host(exit.host, exit.port, timeout)
    try:
        methods = b"\x05\x01\x02" if exit.username else b"\x05\x01\x00"
        connection.sendall(methods)
        version, selected = _recv_exact(connection, 2)
        if version != 5 or selected == 255:
            raise OSError("SOCKS5 authentication method rejected")
        if selected == 2:
            user, credential = exit.username.encode(), exit.credential.encode()
            if len(user) > 255 or len(credential) > 255:
                raise OSError("SOCKS5 credentials too long")
            connection.sendall(
                b"\x01"
                + bytes([len(user)])
                + user
                + bytes([len(credential)])
                + credential
            )
            if _recv_exact(connection, 2)[1] != 0:
                raise OSError("SOCKS5 authentication rejected")
        elif exit.username:
            raise OSError("SOCKS5 proxy skipped required authentication")
        if _is_ipv4(host):
            address = b"\x01" + socket.inet_aton(host)
        else:
            try:
                address = b"\x04" + socket.inet_pton(socket.AF_INET6, host)
            except OSError:
                encoded = host.encode()
                if len(encoded) > 255:
                    raise OSError("SOCKS5 target too long") from None
                address = b"\x03" + bytes([len(encoded)]) + encoded
        connection.sendall(b"\x05\x01\x00" + address + port.to_bytes(2, "big"))
        header = _recv_exact(connection, 4)
        if header[0] != 5 or header[1] != 0:
            raise OSError("SOCKS5 proxy rejected connection")
        if header[3] == 1:
            length = 4
        elif header[3] == 3:
            length = _recv_exact(connection, 1)[0]
        elif header[3] == 4:
            length = 16
        else:
            raise OSError("invalid SOCKS5 response")
        _recv_exact(connection, length + 2)
        connection.settimeout(None)
        return connection
    except Exception:
        connection.close()
        raise


def _read_status(connection: socket.socket) -> int:
    return _read_connect_response(connection)[0]


def _read_connect_response(connection: socket.socket) -> tuple[int, bytes]:
    # Consume the complete response header.  A proxy may return the first
    # tunnel bytes in the same read, so retain bytes after the header block.
    data = bytearray()
    while b"\r\n\r\n" not in data:
        chunk = connection.recv(65536)
        if not chunk:
            raise OSError("proxy response closed before headers")
        data.extend(chunk)
        if len(data) > 256 * 1024:
            raise OSError("proxy response headers too large")
    split = data.index(b"\r\n\r\n") + 4
    header = bytes(data[:split])
    lines = header.decode("iso-8859-1").split("\r\n")
    parts = lines[0].split(" ", 2)
    if len(parts) < 2 or not parts[0].startswith("HTTP/"):
        raise OSError("invalid proxy response status")
    try:
        status = int(parts[1])
    except ValueError as exc:
        raise OSError("invalid proxy response status") from exc
    for line in lines[1:-2]:
        if line and ":" not in line:
            raise OSError("invalid proxy response header")
    return status, bytes(data[split:])


class _BufferedSocket:
    def __init__(self, connection: socket.socket, buffered: bytes) -> None:
        self.connection = connection
        self.buffered = bytearray(buffered)

    def recv(self, size: int, flags: int = 0) -> bytes:
        if self.buffered:
            result = bytes(self.buffered[:size])
            del self.buffered[:size]
            return result
        return (
            self.connection.recv(size, flags) if flags else self.connection.recv(size)
        )

    def sendall(self, data: bytes) -> None:
        self.connection.sendall(data)

    def send(self, data: bytes, flags: int = 0) -> int:
        return (
            self.connection.send(data, flags) if flags else self.connection.send(data)
        )

    def settimeout(self, value: float | None) -> None:
        self.connection.settimeout(value)

    def close(self) -> None:
        self.connection.close()


def _recv_exact(connection: socket.socket, size: int) -> bytes:
    result = bytearray()
    while len(result) < size:
        chunk = connection.recv(size - len(result))
        if not chunk:
            raise OSError("proxy connection closed")
        result.extend(chunk)
    return bytes(result)


def _copy_until_close(source: socket.socket, target: socket.socket) -> None:
    while True:
        data = source.recv(65536)
        if not data:
            return
        target.sendall(data)


def _relay(left: socket.socket, right: socket.socket) -> None:
    sockets = [left, right]
    while sockets:
        readable, _, _ = select.select(sockets, [], [], 60.0)
        if not readable:
            return
        for source in readable:
            destination = right if source is left else left
            data = source.recv(65536)
            if not data:
                return
            destination.sendall(data)


def _is_ipv4(host: str) -> bool:
    try:
        socket.inet_aton(host)
        return True
    except OSError:
        return False