#!/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 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;
        http2 on;
        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_http2(edge_port: int) -> None:
    expected = b'{"code":"OK","message":"success","data":{"status":"ok"},"requestId":"proxyCheck"}\n'
    result = subprocess.run(
        [
            "curl", "--silent", "--show-error", "--fail", "--http2", "--insecure",
            "-H", "X-Request-ID: proxyCheck", "-w", "\\n%{http_version}",
            f"https://127.0.0.1:{edge_port}/health",
        ],
        capture_output=True,
        check=False,
    )
    if result.returncode != 0:
        raise AssertionError(f"HTTP/2 curl failed: {result.stderr.decode(errors='replace')}")
    output = result.stdout
    try:
        body, version = output.rsplit(b"\n", 1)
    except ValueError as error:
        raise AssertionError(f"HTTP/2 curl output missing protocol version: {output!r}") from error
    if version != b"2" or body != expected:
        raise AssertionError(f"expected HTTP/2 response, got version={version!r} body={body!r}")


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 shutil.which("curl") is None:
        raise SystemExit("curl with HTTP/2 support 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", "-i", "--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_http2(edge_port)
            print("PASS Nginx HTTP/2 edge negotiation and HTTP/1.1 native backhaul", flush=True)
            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
            assert backend.stdout is not None
            backend.stdin.write('"snapshot"\n')
            backend.stdin.flush()
            print(f"native after proxy timeout: {backend.stdout.readline().strip()}", flush=True)
            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()