#!/usr/bin/env python3
"""Concurrent load test for pneumatube.

Two targets:

  direct   Talks to a `pneumatube` subprocess directly over LMTP/SMTP,
           bypassing Postfix. Isolates pneumatube's own concurrency
           behaviour -- in particular, SmolWorker (worker.py) serializes
           every smolmail operation onto one thread, so this is where that
           design's throughput ceiling actually shows up.
  postfix  Drives the same load through the real docker-compose.yml rig
           (Postfix -> LMTP -> pneumatube), to confirm nothing changes with
           a real MTA in front. Needs podman (or docker) + compose.

Both targets spin up a throwaway smolmaild, register a gateway account plus
N recipient/sender accounts against it, generate concurrent load, and verify
every expected message actually arrived exactly once -- throughput numbers
are meaningless if messages were silently dropped or duplicated.

Usage:
    uv run python scripts/stress_test.py --target direct
    uv run python scripts/stress_test.py --target postfix
    uv run python scripts/stress_test.py --target both --concurrency 40

Needs a `smolmail` checkout next to this repo (same default pneumatube
itself uses: PNEUMATUBE_SMOLMAIL_REPO, default "../smolmail").
"""
from __future__ import annotations

import argparse
import contextlib
import os
import shutil
import smtplib
import socket
import subprocess
import sys
import tempfile
import threading
import time
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass, field
from pathlib import Path

from aiosmtpd.controller import Controller

sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src"))
from pneumatube import vendor  # noqa: E402

REPO_ROOT = Path(__file__).resolve().parent.parent


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


def wait_for_port(host: str, port: int, timeout: float = 15.0) -> None:
    deadline = time.monotonic() + timeout
    while time.monotonic() < deadline:
        try:
            with socket.create_connection((host, port), timeout=0.5):
                return
        except OSError:
            time.sleep(0.1)
    raise TimeoutError(f"nothing listening on {host}:{port} after {timeout}s")


# --- reporting ----------------------------------------------------------------


@dataclass
class Report:
    label: str
    attempted: int = 0
    succeeded: int = 0
    rejected: int = 0
    errors: int = 0
    latencies: list[float] = field(default_factory=list)

    def record(self, outcome: str, latency: float) -> None:
        self.attempted += 1
        self.latencies.append(latency)
        if outcome == "ok":
            self.succeeded += 1
        elif outcome == "rejected":
            self.rejected += 1
        else:
            self.errors += 1

    def print(self) -> None:
        lat = sorted(self.latencies)
        p50 = lat[len(lat) // 2] if lat else 0.0
        p99 = lat[int(len(lat) * 0.99)] if lat else 0.0
        print(f"-- {self.label} --")
        print(f"  attempted={self.attempted} succeeded={self.succeeded} "
              f"rejected={self.rejected} errors={self.errors}")
        print(f"  latency p50={p50 * 1000:.1f}ms p99={p99 * 1000:.1f}ms "
              f"max={max(lat, default=0) * 1000:.1f}ms")


# --- smolmail test fixtures ----------------------------------------------------


@dataclass
class SmolServer:
    host: str
    port: int
    pubkey_b32: str
    process: subprocess.Popen | None = None


def parse_pubkey(text: str) -> str:
    for line in text.splitlines():
        if "public key" in line:
            return line.strip().rsplit(" ", 1)[-1]
    raise RuntimeError(f"no public key line in output:\n{text}")


def start_smolmaild(smolmail_repo: Path, tmp_dir: Path, rate_connections: int,
                     rate_sends: int) -> SmolServer:
    port = free_port()
    key_path = tmp_dir / "server.key"
    subprocess.run(
        ["uv", "run", str(smolmail_repo / "smolmaild.py"), "keygen", "--key", str(key_path)],
        cwd=smolmail_repo, check=True, capture_output=True, text=True)
    process = subprocess.Popen(
        ["uv", "run", str(smolmail_repo / "smolmaild.py"), "serve",
         "--key", str(key_path), "--db", str(tmp_dir / "server.db"),
         "--host", "127.0.0.1", "--port", str(port),
         # smolmaild's connection/send rate limits are keyed by source IP
         # (SPEC.md S10); every sender in this script shares one IP, so the
         # defaults (120 conn/min, 60 sends/min) would measure the rate
         # limiter, not pneumatube, unless raised for a many-sender test.
         "--rate-connections", str(rate_connections), "--rate-sends", str(rate_sends)],
        cwd=smolmail_repo, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
    pubkey = None
    deadline = time.monotonic() + 15
    while time.monotonic() < deadline:
        line = process.stdout.readline()
        if not line:
            break
        if "server public key" in line:
            pubkey = line.strip().rsplit(" ", 1)[-1]
            break
    if pubkey is None:
        process.kill()
        raise RuntimeError("smolmaild never printed its public key")
    wait_for_port("127.0.0.1", port)
    return SmolServer("127.0.0.1", port, pubkey, process)


@dataclass
class SmolAccount:
    username: str
    key_path: Path
    db_path: Path
    address: str  # "user@host:port"


def make_account(smol, tmp_dir: Path, name: str, server: SmolServer) -> SmolAccount:
    """Keygen + trust + register, reusing smolmail.py's own CLI commands
    directly (the same functions the gateway itself calls) rather than
    reimplementing the protocol or shelling out per account.

    A bare TCP connect (wait_for_port) can succeed slightly before a
    containerized smolmaild can actually complete a Noise handshake, so
    register is retried a few times rather than treating that race as fatal.
    """
    key_path, db_path = tmp_dir / f"{name}.key", tmp_dir / f"{name}.db"
    store = smol.Store(str(db_path))
    smol.cmd_keygen(argparse.Namespace(key=str(key_path), force=True), store)
    smol.cmd_trust(argparse.Namespace(host=f"{server.host}:{server.port}",
                                       key_b32=server.pubkey_b32, force=True), store)
    address = f"{name}@{server.host}:{server.port}"
    register_args = argparse.Namespace(key=str(key_path), address=address, token=None)
    attempts = 20
    for attempt in range(attempts):
        try:
            smol.cmd_register(register_args, store)
            break
        except (smol.SmolError, ConnectionError, OSError):
            if attempt == attempts - 1:
                raise
            time.sleep(1.0)
    return SmolAccount(name, key_path, db_path, address)


def repoint_gateway_account(smol, gateway: SmolAccount, host: str, port: int,
                             pubkey_b32: str) -> None:
    """Point the gateway's own remembered server connection at a different
    host string (e.g. a compose service name) without a real round trip --
    `trust` only ever stores a pin locally (smolmail.py's own cmd_trust)."""
    store = smol.Store(str(gateway.db_path))
    smol.cmd_trust(argparse.Namespace(host=host, key_b32=pubkey_b32, force=True), store)
    store.set_account(smol.Address(gateway.username, host, port))


# --- outbound SMTP sink (stands in for the real relay) -------------------------


class CountingSink:
    def __init__(self) -> None:
        self.lock = threading.Lock()
        self.received: list[tuple[str, tuple[str, ...]]] = []

    async def handle_DATA(self, server, session, envelope):
        with self.lock:
            self.received.append((envelope.mail_from, tuple(envelope.rcpt_tos)))
        return "250 OK"

    def count(self) -> int:
        with self.lock:
            return len(self.received)


# --- load generators (each measures its own latency) ---------------------------


def send_one_lmtp(host: str, port: int, mail_from: str, rcpt: str, body: bytes) -> tuple[str, float]:
    t0 = time.monotonic()
    try:
        with smtplib.LMTP(host, port, timeout=10) as client:
            client.ehlo("stress-test")
            client.mail(mail_from)
            code, _ = client.rcpt(rcpt)
            if code >= 500:
                return "rejected", time.monotonic() - t0
            if code >= 400:
                return "error", time.monotonic() - t0
            client.data(body)
            return "ok", time.monotonic() - t0
    except Exception:
        return "error", time.monotonic() - t0


def send_one_smtp(host: str, port: int, mail_from: str, rcpt: str, body: bytes) -> tuple[str, float]:
    t0 = time.monotonic()
    try:
        with smtplib.SMTP(host, port, timeout=10) as client:
            client.ehlo("stress-test")
            client.mail(mail_from)
            code, _ = client.rcpt(rcpt)
            if code >= 500:
                return "rejected", time.monotonic() - t0
            if code >= 400:
                return "error", time.monotonic() - t0
            client.data(body)
            return "ok", time.monotonic() - t0
    except Exception:
        return "error", time.monotonic() - t0


def send_one_smol(smol, account: SmolAccount, to_address: str, email_to: str,
                   subject: str) -> tuple[str, float]:
    t0 = time.monotonic()
    store = smol.Store(str(account.db_path))
    send_args = argparse.Namespace(
        key=str(account.key_path), address=to_address, body=f"load from {account.username}",
        subject=subject, reply_to=None, header=[f"Email-To: {email_to}"],
        anonymous=True, no_pad=False)
    try:
        smol.cmd_send(send_args, store)
        return "ok", time.monotonic() - t0
    except Exception:
        return "error", time.monotonic() - t0


def run_concurrent(report: Report, concurrency: int, jobs: list) -> list[bool]:
    """jobs: list of (fn, args, expected_valid). Returns per-job valid flags
    in submission order so the caller can cross-check expected vs actual."""
    valid_flags = [valid for _, _, valid in jobs]
    with ThreadPoolExecutor(max_workers=concurrency) as pool:
        start = time.monotonic()
        futures = [pool.submit(fn, *fn_args) for fn, fn_args, _ in jobs]
        for future in futures:
            outcome, latency = future.result()
            report.record(outcome, latency)
        elapsed = time.monotonic() - start
    report.print()
    print(f"  throughput={len(jobs) / elapsed:.1f} msg/s over {elapsed:.1f}s")
    return valid_flags


# --- direct target ---------------------------------------------------------------


def run_direct(args: argparse.Namespace) -> None:
    smolmail_repo = Path(args.smolmail_repo).resolve()
    with tempfile.TemporaryDirectory(prefix="pneumatube-stress-") as tmp:
        tmp_dir = Path(tmp)
        print(f"[direct] starting smolmaild in {tmp_dir}")
        server = start_smolmaild(smolmail_repo, tmp_dir, args.server_rate_connections,
                                  args.server_rate_sends)
        smol = vendor.load(smolmail_repo)

        gateway = make_account(smol, tmp_dir, "gatewayuser", server)
        recipients = [make_account(smol, tmp_dir, f"recipient{i}", server)
                      for i in range(args.recipients)]
        senders = [make_account(smol, tmp_dir, f"sender{i}", server)
                   for i in range(args.concurrency)]
        print(f"[direct] {len(recipients)} recipient account(s), "
              f"{len(senders)} sender account(s) ready")

        sink = CountingSink()
        sink_port = free_port()
        sink_controller = Controller(sink, hostname="127.0.0.1", port=sink_port)
        sink_controller.start()

        listen_port, lookup_port = free_port(), free_port()
        env = {**os.environ,
               "PNEUMATUBE_SMOLMAIL_REPO": str(smolmail_repo),
               "PNEUMATUBE_SMOLMAIL_KEY": str(gateway.key_path),
               "PNEUMATUBE_SMOLMAIL_DB": str(gateway.db_path),
               "PNEUMATUBE_SMOLMAIL_DOMAIN": f"{server.host}:{server.port}",
               "PNEUMATUBE_EMAIL_DOMAIN": "gateway.example",
               "PNEUMATUBE_FETCH_INTERVAL": "2",
               # Throughput, not the accept-gate features, is this script's
               # concern -- the throwaway recipient/sender accounts here
               # never go through the accept dance with the gateway.
               "PNEUMATUBE_REQUIRE_ACCEPTED_SENDER": "0",
               "PNEUMATUBE_REQUIRE_ACCEPTED_RECIPIENT": "0",
               "PNEUMATUBE_LISTEN_PROTOCOL": "lmtp",
               "PNEUMATUBE_LISTEN_HOST": "127.0.0.1",
               "PNEUMATUBE_LISTEN_PORT": str(listen_port),
               "PNEUMATUBE_LOOKUP_HOST": "127.0.0.1",
               "PNEUMATUBE_LOOKUP_PORT": str(lookup_port),
               "PNEUMATUBE_SMTP_RELAY_HOST": "127.0.0.1",
               "PNEUMATUBE_SMTP_RELAY_PORT": str(sink_port),
               "PNEUMATUBE_SMTP_RELAY_STARTTLS": "0",
               "PNEUMATUBE_GATEWAY_FROM_ADDRESS": "gateway@gateway.example"}
        print("[direct] starting pneumatube")
        proc = subprocess.Popen(["uv", "run", "pneumatube"], cwd=REPO_ROOT, env=env,
                                 stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
        try:
            wait_for_port("127.0.0.1", listen_port)
            wait_for_port("127.0.0.1", lookup_port)

            # Phase A: email -> smolmail, concurrent LMTP, mixed valid/invalid.
            total = args.concurrency * args.messages
            jobs = []
            for i in range(total):
                valid = i % 3 != 0  # 2/3 valid, 1/3 unknown-user
                rcpt = (f"{recipients[i % len(recipients)].username}@gateway.example"
                        if valid else "nobody-at-all@gateway.example")
                body = f"Subject: load {i}\r\n\r\nmessage {i}\r\n".encode()
                jobs.append((send_one_lmtp, ("127.0.0.1", listen_port,
                                              "stress@sender.example", rcpt, body), valid))
            report_a = Report("email -> smolmail (LMTP, direct)")
            valid_flags = run_concurrent(report_a, args.concurrency, jobs)
            expected_ok, expected_rejected = sum(valid_flags), valid_flags.count(False)
            ok = report_a.succeeded == expected_ok and report_a.rejected == expected_rejected
            print(f"  correctness: {'PASS' if ok else 'FAIL'} "
                  f"(expected ok={expected_ok} rejected={expected_rejected})")

            # Phase B: smolmail -> email, concurrent SEND, then drain via the
            # gateway's own poll loop (already running in the subprocess).
            jobs = [(send_one_smol, (smol, senders[i % len(senders)], gateway.address,
                                      f"dest{i}@example.org", f"load {i}"), True)
                    for i in range(total)]
            report_b = Report("smolmail -> email (SEND + poll drain)")
            run_concurrent(report_b, args.concurrency, jobs)

            drain_deadline = time.monotonic() + max(30.0, total * 0.05)
            while sink.count() < report_b.succeeded and time.monotonic() < drain_deadline:
                time.sleep(0.5)
            drained = sink.count()
            print(f"  drained {drained}/{report_b.succeeded} into the outbound sink "
                  f"({'PASS' if drained == report_b.succeeded else 'FAIL: lost or stuck messages'})")
        finally:
            proc.terminate()
            with contextlib.suppress(Exception):
                proc.wait(timeout=5)
            sink_controller.stop()
            server.process.terminate()
            with contextlib.suppress(Exception):
                server.process.wait(timeout=5)


# --- postfix target ---------------------------------------------------------------


def podman(*cmd: str, **kwargs) -> subprocess.CompletedProcess:
    return subprocess.run(["podman", *cmd], text=True, capture_output=True, **kwargs)


def run_postfix(args: argparse.Namespace) -> None:
    if shutil.which("podman") is None:
        print("[postfix] podman not found; skipping", file=sys.stderr)
        return

    smolmail_repo = Path(args.smolmail_repo).resolve()
    smol = vendor.load(smolmail_repo)
    network = "pneumatube-stress-net"
    env = {**os.environ,
           "PNEUMATUBE_SMOLMAIL_DOMAIN": "smolmail-dev:1961",
           "PNEUMATUBE_EMAIL_DOMAIN": "gateway.example",
           # Throughput, not the accept-gate features, is this script's
           # concern -- the throwaway recipient accounts here never go
           # through the accept dance with the gateway.
           "PNEUMATUBE_REQUIRE_ACCEPTED_SENDER": "0",
           "PNEUMATUBE_REQUIRE_ACCEPTED_RECIPIENT": "0",
           "PNEUMATUBE_SMTP_RELAY_HOST": "smolmail-dev",
           "PNEUMATUBE_SMTP_RELAY_PORT": "2626",
           "PNEUMATUBE_SMTP_RELAY_STARTTLS": "0",
           "PNEUMATUBE_GATEWAY_FROM_ADDRESS": "gateway@gateway.example"}

    with tempfile.TemporaryDirectory(prefix="pneumatube-stress-pf-") as tmp:
        tmp_dir = Path(tmp)
        data_dir = REPO_ROOT / "data"
        data_dir.mkdir(exist_ok=True)
        try:
            print("[postfix] building images")
            subprocess.run(["podman", "compose", "build"], cwd=REPO_ROOT, env=env, check=True)

            print("[postfix] starting throwaway smolmaild-dev container")
            podman("network", "create", network)
            server_dir = tmp_dir / "server"
            server_dir.mkdir()
            keygen = podman("run", "--rm", "-v", f"{smolmail_repo}:/smolmail:ro",
                             "-v", f"{server_dir}:/data",
                             "localhost/pneumatube_pneumatube:latest",
                             "uv", "run", "/smolmail/smolmaild.py", "keygen", "--key", "/data/server.key")
            keygen.check_returncode()
            pubkey = parse_pubkey(keygen.stdout)

            host_port = free_port()
            podman("run", "-d", "--name", "pneumatube-stress-smolmaild",
                   "--network", network, "--network-alias", "smolmail-dev",
                   "-p", f"{host_port}:1961",
                   "-v", f"{smolmail_repo}:/smolmail:ro", "-v", f"{server_dir}:/data",
                   "localhost/pneumatube_pneumatube:latest",
                   "uv", "run", "/smolmail/smolmaild.py", "serve",
                   "--key", "/data/server.key", "--db", "/data/server.db",
                   "--host", "0.0.0.0", "--port", "1961",
                   "--rate-connections", str(args.server_rate_connections),
                   "--rate-sends", str(args.server_rate_sends)).check_returncode()
            wait_for_port("127.0.0.1", host_port)

            server = SmolServer("127.0.0.1", host_port, pubkey)
            gateway = make_account(smol, tmp_dir, "gatewayuser", server)
            recipients = [make_account(smol, tmp_dir, f"recipient{i}", server)
                          for i in range(args.recipients)]
            repoint_gateway_account(smol, gateway, "smolmail-dev", 1961, pubkey)
            (data_dir / "gateway.key").write_bytes(gateway.key_path.read_bytes())
            (data_dir / "gateway.db").write_bytes(gateway.db_path.read_bytes())

            print("[postfix] starting pneumatube + postfix")
            subprocess.run(["podman", "compose", "up", "-d"], cwd=REPO_ROOT, env=env, check=True)
            # Join the compose-started containers to the network carrying
            # smolmail-dev (they start on compose's own default network,
            # which smolmail-dev was never part of) -- not the other way
            # around, which would just reconnect smolmail-dev to a network
            # it's already on and leave pneumatube unable to resolve it.
            podman("network", "connect", network, "pneumatube_pneumatube_1")
            podman("network", "connect", network, "pneumatube_postfix_1")
            wait_for_port("127.0.0.1", 2525, timeout=60.0)
            # Port 25 accepting doesn't mean postconf's chroot/transport_maps
            # setup (entrypoint.sh) has finished; give it a moment to settle.
            time.sleep(2)

            total = args.concurrency * args.messages
            jobs = []
            for i in range(total):
                valid = i % 3 != 0
                rcpt = (f"{recipients[i % len(recipients)].username}@gateway.example"
                        if valid else "nobody-at-all@gateway.example")
                body = f"Subject: pf-load {i}\r\n\r\nmessage {i}\r\n".encode()
                jobs.append((send_one_smtp, ("127.0.0.1", 2525,
                                              "stress@sender.example", rcpt, body), valid))
            report = Report("email -> smolmail (through Postfix)")
            valid_flags = run_concurrent(report, args.concurrency, jobs)
            expected_ok, expected_rejected = sum(valid_flags), valid_flags.count(False)
            print(f"  correctness (live reject, not bounce): "
                  f"{'PASS' if report.rejected == expected_rejected else 'FAIL'} "
                  f"(expected ok={expected_ok} rejected={expected_rejected})")

            deadline = time.monotonic() + max(30.0, total * 0.1)
            delivered = 0
            while time.monotonic() < deadline:
                delivered = 0
                for r in recipients:
                    store = smol.Store(str(r.db_path))
                    smol.cmd_fetch(argparse.Namespace(key=str(r.key_path), db=str(r.db_path),
                                                        reset=False, keep=False), store)
                    delivered += len(store.mail("all"))
                if delivered >= expected_ok:
                    break
                time.sleep(1)
            print(f"  delivered {delivered}/{expected_ok} across recipient inboxes "
                  f"({'PASS' if delivered == expected_ok else 'FAIL: lost or stuck messages'})")
        finally:
            print("[postfix] tearing down")
            subprocess.run(["podman", "compose", "down"], cwd=REPO_ROOT, env=env)
            podman("rm", "-f", "pneumatube-stress-smolmaild")
            podman("network", "rm", network)
            with contextlib.suppress(Exception):
                (data_dir / "gateway.key").unlink()
                (data_dir / "gateway.db").unlink()
                data_dir.rmdir()


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__,
                                      formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--target", choices=["direct", "postfix", "both"], default="direct")
    parser.add_argument("--concurrency", type=int, default=20,
                         help="concurrent sender workers per phase")
    parser.add_argument("--messages", type=int, default=10,
                         help="messages sent by each worker")
    parser.add_argument("--recipients", type=int, default=10,
                         help="number of pre-registered valid recipient accounts")
    parser.add_argument("--smolmail-repo", default="../smolmail")
    parser.add_argument("--server-rate-connections", type=int, default=6000,
                         help="smolmaild --rate-connections for the throwaway test server "
                              "(default 120/min would rate-limit this whole test, since every "
                              "sender shares one source IP; pass 120 to test against that "
                              "realistic default instead)")
    parser.add_argument("--server-rate-sends", type=int, default=6000,
                         help="smolmaild --rate-sends for the throwaway test server "
                              "(see --server-rate-connections; realistic default is 60)")
    args = parser.parse_args()

    if args.target in ("direct", "both"):
        run_direct(args)
    if args.target in ("postfix", "both"):
        run_postfix(args)


if __name__ == "__main__":
    main()