#!/usr/bin/env python3
"""Verify an isolated Nginx -> native TLS reverse-proxy deployment on Linux."""

import argparse
import http.client
import json
from pathlib import Path
import selectors
import shutil
import socket
import ssl
import subprocess
import sys
import tempfile
import time


def free_port() -> int:
    with socket.socket() as listener:
        listener.bind(("127.0.0.1", 0))
        return listener.getsockname()[1]


def wait_for_native(process: subprocess.Popen[str], timeout: float) -> int:
    selector = selectors.DefaultSelector()
    assert process.stdout is not None
    selector.register(process.stdout, selectors.EVENT_READ)
    deadline = time.monotonic() + timeout
    try:
        while time.monotonic() < deadline:
            if process.poll() is not None:
                raise RuntimeError(f"native server exited with code {process.returncode}")
            for _, _ in selector.select(timeout=0.5):
                line = process.stdout.readline()
                if not line:
                    continue
                try:
                    value = json.loads(line)
                except json.JSONDecodeError:
                    continue
                address = value.get("address", "")
                if address:
                    return int(address.rsplit(":", 1)[1])
        raise RuntimeError("native TLS server did not announce a listener")
    finally:
        selector.close()


def write_edge_certificate(directory: Path) -> None:
    subprocess.run(
        [
            "openssl", "req", "-x509", "-newkey", "rsa:2048", "-nodes",
            "-keyout", str(directory / "key.pem"),
            "-out", str(directory / "cert.pem"),
            "-days", "1", "-subj", "/CN=localhost",
        ],
        check=True,
        stdout=subprocess.DEVNULL,
        stderr=subprocess.DEVNULL,
    )


def write_nginx_config(path: Path, backend_port: int, edge_port: int) -> None:
    path.write_text(
        f"""events {{
    worker_connections 512;
}}

http {{
    upstream go_omni_native {{
        server 127.0.0.1:{backend_port};
        keepalive 64;
    }}

    server {{
        listen 127.0.0.1:{edge_port} ssl;
        ssl_certificate /etc/nginx/certs/cert.pem;
        ssl_certificate_key /etc/nginx/certs/key.pem;

        proxy_http_version 1.1;
        proxy_set_header Connection "";
        proxy_set_header Host $host;
        proxy_set_header X-Request-ID $http_x_request_id;
        proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for;
        proxy_connect_timeout 1s;
        proxy_read_timeout 300ms;
        proxy_send_timeout 1s;

        location / {{
            proxy_pass https://go_omni_native;
            proxy_ssl_verify off;
            proxy_ssl_server_name off;
        }}
    }}
}}
"""
    )


def request(connection: http.client.HTTPSConnection, path: str) -> tuple[int, dict[str, str], bytes]:
    connection.request("GET", path, headers={"X-Request-ID": "proxyCheck"})
    response = connection.getresponse()
    body = response.read()
    return response.status, dict(response.headers), body


def header(headers: dict[str, str], name: str) -> str:
    for key, value in headers.items():
        if key.lower() == name.lower():
            return value
    return ""


def verify_keep_alive(edge_port: int) -> None:
    context = ssl._create_unverified_context()
    connection = http.client.HTTPSConnection("127.0.0.1", edge_port, context=context, timeout=3)
    expected = b'{"code":"OK","message":"success","data":{"status":"ok"},"requestId":"proxyCheck"}\n'
    try:
        for attempt in range(5):
            status, headers, body = request(connection, "/health")
            if status != 200 or header(headers, "X-Request-ID") != "proxyCheck" or body != expected:
                raise AssertionError(f"unexpected proxied response: status={status} body={body!r}")
            if attempt < 4 and getattr(connection, "sock", None) is None:
                raise AssertionError("Nginx did not keep the TLS client connection alive")
    finally:
        connection.close()


def wait_for_proxy(edge_port: int, timeout: float) -> None:
    deadline = time.monotonic() + timeout
    last_error = "proxy did not return a healthy response"
    while time.monotonic() < deadline:
        connection = None
        try:
            context = ssl._create_unverified_context()
            connection = http.client.HTTPSConnection("127.0.0.1", edge_port, context=context, timeout=2)
            status, headers, body = request(connection, "/health")
            if status == 200 and header(headers, "X-Request-ID") == "proxyCheck" and body.endswith(b"\n"):
                return
            last_error = f"status={status} body={body!r}"
        except (OSError, http.client.HTTPException) as error:
            last_error = str(error)
        finally:
            if connection is not None:
                connection.close()
        time.sleep(0.2)
    raise RuntimeError(f"Nginx proxy did not become healthy: {last_error}")


def verify_timeout(edge_port: int) -> None:
    context = ssl._create_unverified_context()
    connection = http.client.HTTPSConnection("127.0.0.1", edge_port, context=context, timeout=3)
    try:
        status, _, _ = request(connection, "/slow")
        if status != 504:
            raise AssertionError(f"expected Nginx upstream timeout 504, got {status}")
    finally:
        connection.close()


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--image", default="golang:1.25")
    parser.add_argument("--nginx-image", default="nginx:1.27-alpine")
    args = parser.parse_args()
    if shutil.which("docker") is None:
        raise SystemExit("docker is required for proxy verification")
    if shutil.which("openssl") is None:
        raise SystemExit("openssl is required for proxy verification")
    if not Path("/proc").exists():
        raise SystemExit("proxy verification must run on Linux")

    root = Path(__file__).resolve().parent.parent
    backend_name = "go-omni-proxy-native"
    nginx_name = "go-omni-proxy-nginx"
    backend = None
    nginx = None
    with tempfile.TemporaryDirectory(prefix="go-omni-proxy-") as temporary:
        directory = Path(temporary)
        edge_port = free_port()
        write_edge_certificate(directory)
        backend = subprocess.Popen(
            [
                "docker", "run", "--rm", "--pull=missing", "--name", backend_name,
                "--network", "host", "--cpus=1", "--memory=512m", "--pids-limit=512",
                "-v", f"{root}:/workspace", "-w", "/workspace", "-e", "GOWORK=off",
                args.image, "sh", "-c", "go run ./internal/perfcheck --server native --tls",
            ],
            stdin=subprocess.PIPE,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            bufsize=1,
        )
        try:
            backend_port = wait_for_native(backend, 120)
            write_nginx_config(directory / "nginx.conf", backend_port, edge_port)
            print(f"proxy backend=127.0.0.1:{backend_port} edge=127.0.0.1:{edge_port}", flush=True)
            nginx = subprocess.Popen(
                [
                    "docker", "run", "--rm", "--pull=missing", "--name", nginx_name,
                    "--network", "host",
                    "-v", f"{directory / 'nginx.conf'}:/etc/nginx/nginx.conf:ro",
                    "-v", f"{directory}:/etc/nginx/certs:ro",
                    args.nginx_image, "nginx", "-g", "daemon off;",
                ],
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True,
            )
            deadline = time.monotonic() + 30
            while time.monotonic() < deadline:
                if nginx.poll() is not None:
                    raise RuntimeError(f"Nginx exited with code {nginx.returncode}")
                try:
                    with socket.create_connection(("127.0.0.1", edge_port), timeout=0.5):
                        break
                except OSError:
                    time.sleep(0.2)
            else:
                raise RuntimeError("Nginx did not start listening")

            try:
                wait_for_proxy(edge_port, 30)
            except Exception as error:
                nginx_logs = subprocess.run(
                    ["docker", "logs", nginx_name], capture_output=True, text=True, check=False
                ).stdout[-2000:]
                backend_logs = subprocess.run(
                    ["docker", "logs", backend_name], capture_output=True, text=True, check=False
                ).stdout[-2000:]
                raise RuntimeError(
                    f"{error}; nginx logs={nginx_logs!r}; backend logs={backend_logs!r}"
                ) from error
            verify_keep_alive(edge_port)
            verify_timeout(edge_port)
            print("PASS Nginx TLS proxy, Keep-Alive and upstream timeout", flush=True)

            assert backend.stdin is not None
            backend.stdin.write('"stop"\n')
            backend.stdin.flush()
            exit_code = backend.wait(timeout=5)
            if exit_code != 0:
                details = backend.stderr.read().strip() if backend.stderr is not None else ""
                raise AssertionError(
                    f"native server did not shut down cleanly behind Nginx: exit={exit_code} {details}"
                )
            print("PASS native graceful shutdown behind proxy", flush=True)
        finally:
            if nginx is not None and nginx.poll() is None:
                nginx.terminate()
                try:
                    nginx.wait(timeout=5)
                except subprocess.TimeoutExpired:
                    nginx.kill()
                    nginx.wait(timeout=5)
            if backend is not None and backend.poll() is None:
                if backend.stdin is not None:
                    try:
                        backend.stdin.write('"stop"\n')
                        backend.stdin.flush()
                    except OSError:
                        pass
                backend.terminate()
                try:
                    backend.wait(timeout=5)
                except subprocess.TimeoutExpired:
                    backend.kill()
                    backend.wait(timeout=5)


if __name__ == "__main__":
    main()