#!/usr/bin/env python3
"""Bounded HTTP CONNECT proxy used behind the Mistral SSH tunnel."""

from __future__ import annotations

import os
import selectors
import socket
import socketserver
import threading
import time
from typing import Any


LISTEN_HOST = "127.0.0.1"
LISTEN_PORT = 18081
MAX_HEADER = 65536
TIMEOUT_SEC = 30
IDLE_TIMEOUT_SEC = 900
BUFFER_SIZE = 65536


def read_env_int(name: str, default: int) -> int:
    raw = os.getenv(name, "").strip()
    if not raw:
        return default
    try:
        return int(raw)
    except ValueError:
        return default


class Handler(socketserver.BaseRequestHandler):
    def handle(self) -> None:
        self.request.settimeout(TIMEOUT_SEC)
        data = b""
        try:
            while b"\r\n\r\n" not in data and len(data) < MAX_HEADER:
                chunk = self.request.recv(4096)
                if not chunk:
                    return
                data += chunk
        except OSError:
            return

        header, _, rest = data.partition(b"\r\n\r\n")
        first = header.split(b"\r\n", 1)[0].decode("latin1", "replace")
        parts = first.split()
        if len(parts) < 3 or parts[0].upper() != "CONNECT":
            self.send_response(b"HTTP/1.1 405 Method Not Allowed\r\nConnection: close\r\n\r\n")
            return

        target = parts[1]
        if ":" not in target:
            self.send_response(b"HTTP/1.1 400 Bad Request\r\nConnection: close\r\n\r\n")
            return

        host, port_text = target.rsplit(":", 1)
        try:
            port = int(port_text)
        except ValueError:
            self.send_response(b"HTTP/1.1 400 Bad Request\r\nConnection: close\r\n\r\n")
            return

        try:
            upstream = socket.create_connection((host, port), timeout=TIMEOUT_SEC)
        except OSError:
            self.send_response(b"HTTP/1.1 502 Bad Gateway\r\nConnection: close\r\n\r\n")
            return

        with upstream:
            try:
                self.request.sendall(
                    b"HTTP/1.1 200 Connection Established\r\nProxy-Agent: yunwo-connect\r\n\r\n"
                )
                if rest:
                    upstream.sendall(rest)
            except OSError:
                return

            upstream.setblocking(False)
            self.request.setblocking(False)
            with selectors.DefaultSelector() as selector:
                selector.register(self.request, selectors.EVENT_READ, upstream)
                selector.register(upstream, selectors.EVENT_READ, self.request)
                idle_deadline = time.monotonic() + IDLE_TIMEOUT_SEC
                while time.monotonic() < idle_deadline:
                    events = selector.select(timeout=TIMEOUT_SEC)
                    if not events:
                        continue
                    for key, _ in events:
                        source = key.fileobj
                        destination = key.data
                        try:
                            chunk = source.recv(BUFFER_SIZE)
                        except BlockingIOError:
                            continue
                        except OSError:
                            return
                        if not chunk:
                            return
                        idle_deadline = time.monotonic() + IDLE_TIMEOUT_SEC
                        try:
                            destination.sendall(chunk)
                        except OSError:
                            return

    def send_response(self, response: bytes) -> None:
        try:
            self.request.sendall(response)
        except OSError:
            pass


class BoundedThreadingTCPServer(socketserver.ThreadingMixIn, socketserver.TCPServer):
    allow_reuse_address = True
    daemon_threads = True
    block_on_close = False
    request_queue_size = 64

    def __init__(
        self,
        server_address: tuple[str, int],
        request_handler: type[socketserver.BaseRequestHandler],
        *,
        max_workers: int,
    ) -> None:
        self.max_workers = max(1, max_workers)
        self._worker_semaphore = threading.BoundedSemaphore(self.max_workers)
        super().__init__(server_address, request_handler)

    def process_request(self, request: Any, client_address: tuple[str, int]) -> None:
        if not self._worker_semaphore.acquire(blocking=False):
            self.send_overloaded_response(request)
            self.shutdown_request(request)
            return
        try:
            super().process_request(request, client_address)
        except Exception:
            self._worker_semaphore.release()
            raise

    def process_request_thread(self, request: Any, client_address: tuple[str, int]) -> None:
        try:
            super().process_request_thread(request, client_address)
        finally:
            self._worker_semaphore.release()

    @staticmethod
    def send_overloaded_response(request: Any) -> None:
        try:
            request.sendall(
                b"HTTP/1.1 503 Service Unavailable\r\n"
                b"Connection: close\r\n"
                b"Content-Length: 0\r\n\r\n"
            )
        except OSError:
            pass


if __name__ == "__main__":
    max_workers = max(1, read_env_int("YUNWO_CONNECT_PROXY_MAX_CONCURRENCY", 32))
    with BoundedThreadingTCPServer(
        (LISTEN_HOST, LISTEN_PORT),
        Handler,
        max_workers=max_workers,
    ) as server:
        print(
            f"yunwo HTTP CONNECT proxy listening on {LISTEN_HOST}:{LISTEN_PORT} "
            f"max_workers={max_workers}",
            flush=True,
        )
        server.serve_forever()