"""Per-bottle TCP forwarder for smolmachines sidecar VMs.

smolVM currently publishes guest ports on host loopback as
``HOST:GUEST`` port pairs, without an address-scoped bind. The agent
VM's TSI allowlist is IP-only, so bot-bottle exposes a second, stricter
surface: listeners bound only to the bottle's allocated host address,
forwarding to the raw smolVM-published loopback ports.
"""

from __future__ import annotations

import argparse
import json
import os
import selectors
import signal
import socket
import subprocess
import sys
import threading
from dataclasses import asdict, dataclass
from typing import Any, Iterable, Sequence

from ...log import die, warn


_BUFFER_SIZE = 64 * 1024
_LOOPBACK_TARGETS = {"127.0.0.1", "::1"}
_FORBIDDEN_LISTEN_HOSTS = {"", "0.0.0.0", "::", "127.0.0.1", "::1"}


@dataclass(frozen=True)
class ForwardSpec:
    label: str
    listen_host: str
    listen_port: int
    target_host: str
    target_port: int


@dataclass(frozen=True)
class ForwarderHandle:
    process: Any
    forwards: tuple[ForwardSpec, ...]


def start_forwarder(specs: Sequence[ForwardSpec]) -> ForwarderHandle:
    """Start one helper process for a bottle and wait until all
    listeners are bound. ``listen_port`` may be 0; the returned specs
    contain the kernel-assigned ports printed by the helper."""
    if not specs:
        return ForwarderHandle(process=_CompletedProcess(), forwards=())
    _validate_specs(specs)
    argv = [
        sys.executable,
        "-c",
        "from bot_bottle.backend.smolmachines.port_forward import main; "
        "raise SystemExit(main())",
        "--spec-json",
        json.dumps([asdict(s) for s in specs], separators=(",", ":")),
    ]
    proc = subprocess.Popen(
        argv,
        stdin=subprocess.DEVNULL,
        stdout=subprocess.PIPE,
        stderr=None,
        text=True,
        env=os.environ,
    )
    assert proc.stdout is not None
    line = proc.stdout.readline()
    proc.stdout.close()
    if not line:
        proc.terminate()
        proc.wait(timeout=5)
        die("smolmachines forwarder failed before reporting readiness")
    try:
        payload = json.loads(line)
        bound = tuple(ForwardSpec(**item) for item in payload["forwards"])
    except (KeyError, TypeError, ValueError, json.JSONDecodeError) as exc:
        proc.terminate()
        proc.wait(timeout=5)
        die(f"smolmachines forwarder returned invalid readiness payload: {exc}")
    return ForwarderHandle(process=proc, forwards=bound)


def stop_forwarder(handle: ForwarderHandle) -> None:
    proc = handle.process
    if proc.poll() is not None:
        return
    proc.terminate()
    try:
        proc.wait(timeout=5)
    except subprocess.TimeoutExpired:
        proc.kill()
        proc.wait(timeout=5)


def _validate_specs(specs: Iterable[ForwardSpec]) -> None:
    for spec in specs:
        if spec.listen_host in _FORBIDDEN_LISTEN_HOSTS:
            die(f"refusing unsafe smolmachines forwarder bind: {spec.listen_host!r}")
        if spec.target_host not in _LOOPBACK_TARGETS:
            die(f"refusing non-loopback smolmachines forwarder target: {spec.target_host!r}")
        if not 0 <= spec.listen_port <= 65535:
            die(f"invalid smolmachines forwarder listen port: {spec.listen_port}")
        if not 1 <= spec.target_port <= 65535:
            die(f"invalid smolmachines forwarder target port: {spec.target_port}")


class _CompletedProcess:
    """Popen-like no-op used when a bottle has no forwards."""

    def poll(self) -> int:
        return 0

    def terminate(self) -> None:
        return None

    def kill(self) -> None:
        return None

    def wait(self, timeout: float | None = None) -> int:
        del timeout
        return 0


class _Forwarder:
    def __init__(self, specs: Sequence[ForwardSpec]):
        _validate_specs(specs)
        self.specs = tuple(specs)
        self.selector = selectors.DefaultSelector()
        self.stop_event = threading.Event()
        self.listeners: list[tuple[socket.socket, ForwardSpec]] = []
        self.bound: list[ForwardSpec] = []

    def start(self) -> None:
        for spec in self.specs:
            listener = socket.create_server(
                (spec.listen_host, spec.listen_port),
                reuse_port=False,
                backlog=64,
            )
            listener.setblocking(False)
            host, port = listener.getsockname()[:2]
            bound = ForwardSpec(
                label=spec.label,
                listen_host=str(host),
                listen_port=int(port),
                target_host=spec.target_host,
                target_port=spec.target_port,
            )
            self.listeners.append((listener, bound))
            self.bound.append(bound)
            self.selector.register(listener, selectors.EVENT_READ, bound)

    def serve(self) -> None:
        while not self.stop_event.is_set():
            try:
                events = self.selector.select(timeout=0.25)
            except OSError:
                return
            for key, _ in events:
                listener = key.fileobj
                if not isinstance(listener, socket.socket):
                    continue
                spec = key.data
                try:
                    client, _ = listener.accept()
                except OSError:
                    continue
                threading.Thread(
                    target=self._handle_client,
                    args=(client, spec),
                    daemon=True,
                ).start()

    def stop(self) -> None:
        self.stop_event.set()
        for listener, _ in self.listeners:
            try:
                self.selector.unregister(listener)
            except Exception:  # noqa: BLE001 - best effort shutdown
                pass
            listener.close()
        self.selector.close()

    def _handle_client(self, client: socket.socket, spec: ForwardSpec) -> None:
        with client:
            try:
                target = socket.create_connection((spec.target_host, spec.target_port))
            except OSError as exc:
                warn(f"smolmachines forwarder {spec.label}: target connect failed: {exc}")
                return
            with target:
                left = threading.Thread(target=_pipe, args=(client, target), daemon=True)
                right = threading.Thread(target=_pipe, args=(target, client), daemon=True)
                left.start()
                right.start()
                left.join()
                right.join()


def _pipe(src: socket.socket, dst: socket.socket) -> None:
    while True:
        try:
            data = src.recv(_BUFFER_SIZE)
        except OSError:
            break
        if not data:
            break
        try:
            dst.sendall(data)
        except OSError:
            break
    try:
        dst.shutdown(socket.SHUT_WR)
    except OSError:
        pass


def main(argv: Sequence[str] | None = None) -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--spec-json", required=True)
    ns = parser.parse_args(argv)
    specs = tuple(ForwardSpec(**item) for item in json.loads(ns.spec_json))
    forwarder = _Forwarder(specs)

    def request_stop(signum: int, _frame: object) -> None:
        del signum
        forwarder.stop()

    signal.signal(signal.SIGTERM, request_stop)
    signal.signal(signal.SIGINT, request_stop)
    forwarder.start()
    print(json.dumps({"forwards": [asdict(s) for s in forwarder.bound]}), flush=True)
    try:
        forwarder.serve()
    finally:
        forwarder.stop()
    return 0


if __name__ == "__main__":
    raise SystemExit(main())