"""Localhost-only TCP port proxy for published workspace services."""

from __future__ import annotations

import argparse
import json
import selectors
import signal
import socket
import socketserver
import sys
import threading
from pathlib import Path

DEFAULT_PUBLISHED_PORT_HOST = "127.0.0.1"


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

    def __init__(self, server_address: tuple[str, int], target_address: tuple[str, int]) -> None:
        super().__init__(server_address, _ProxyHandler)
        self.target_address = target_address


class _ProxyHandler(socketserver.BaseRequestHandler):
    def handle(self) -> None:
        server = self.server
        if not isinstance(server, _ProxyServer):
            raise RuntimeError("proxy server is invalid")
        try:
            upstream = socket.create_connection(server.target_address, timeout=5)
        except OSError:
            return
        with upstream:
            self.request.setblocking(False)
            upstream.setblocking(False)
            selector = selectors.DefaultSelector()
            try:
                selector.register(self.request, selectors.EVENT_READ, upstream)
                selector.register(upstream, selectors.EVENT_READ, self.request)
                while True:
                    events = selector.select()
                    if not events:
                        continue
                    for key, _ in events:
                        source = key.fileobj
                        target = key.data
                        if not isinstance(source, socket.socket) or not isinstance(
                            target, socket.socket
                        ):
                            continue
                        try:
                            chunk = source.recv(65536)
                        except OSError:
                            return
                        if not chunk:
                            return
                        try:
                            target.sendall(chunk)
                        except OSError:
                            return
            finally:
                selector.close()


def _build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="Run a localhost-only TCP port proxy.")
    parser.add_argument("--listen-host", required=True)
    parser.add_argument("--listen-port", type=int, required=True)
    parser.add_argument("--target-host", required=True)
    parser.add_argument("--target-port", type=int, required=True)
    parser.add_argument("--ready-file", required=True)
    return parser


def main(argv: list[str] | None = None) -> int:
    args = _build_parser().parse_args(argv)
    ready_file = Path(args.ready_file)
    ready_file.parent.mkdir(parents=True, exist_ok=True)
    server = _ProxyServer(
        (str(args.listen_host), int(args.listen_port)),
        (str(args.target_host), int(args.target_port)),
    )
    actual_host = str(server.server_address[0])
    actual_port = int(server.server_address[1])
    ready_file.write_text(
        json.dumps(
            {
                "host": actual_host,
                "host_port": actual_port,
                "target_host": args.target_host,
                "target_port": int(args.target_port),
                "protocol": "tcp",
            },
            indent=2,
            sort_keys=True,
        ),
        encoding="utf-8",
    )

    def _shutdown(_: int, __: object) -> None:
        threading.Thread(target=server.shutdown, daemon=True).start()

    signal.signal(signal.SIGTERM, _shutdown)
    signal.signal(signal.SIGINT, _shutdown)
    try:
        server.serve_forever(poll_interval=0.2)
    finally:
        server.server_close()
    return 0


if __name__ == "__main__":
    raise SystemExit(main(sys.argv[1:]))
