"""R-T1-9D (survivability §3.7, Tier-1 #9, S5) — per-tenant rate limit
(token bucket).

The lease / connection / slowloris caps bound resource HOLDING; this
bounds op THROUGHPUT per tenant so a flood (op storm) can't saturate a
worker or burn the broker's CPU. Each control-path op costs one token;
tokens refill at a steady rate up to a burst ceiling; an op with no
token is shed with `Err.Overloaded`.

Production defaults are deliberately generous (a real workload never
trips them); these tests spawn the broker with a tiny burst via
`MGF_HRB_RL_BURST` / `MGF_HRB_RL_REFILL_PER_SEC` so the shed — and the
refill recovery — are observable deterministically.
"""

from __future__ import annotations

import json
import os
import socket
import subprocess
import time
from pathlib import Path

import pytest

from tests import _broker

pytestmark = pytest.mark.contract

_REPO = Path(__file__).resolve().parents[2]
_SERVER = _REPO / "hrb-code" / "zig-out" / "bin" / "mgf-hrb-server"
_SOCK = _broker.default_sock_path()  # per-xdist-worker socket


def _spawn(env_overrides: dict[str, str]) -> subprocess.Popen[bytes]:
    if Path(_SOCK).exists():
        Path(_SOCK).unlink()
    env = {
        "PATH": os.environ.get("PATH", "/usr/bin:/bin"),
        "MGF_HRB_SOCK": _SOCK,
        **env_overrides,
    }
    proc = subprocess.Popen(  # noqa: S603
        [str(_SERVER)],
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        env=env,
    )
    # Wait until the broker ACCEPTS connections, not merely until the socket
    # file exists: a Unix socket appears at bind() but refuses connects until
    # listen(), so an existence check races (`_connect_hello` → ECONNREFUSED
    # under contention). Reuse the shared connect-probe helper rather than
    # re-deriving the (previously buggy) wait here.
    try:
        _broker.wait_for_socket(_SOCK, 5.0)
    except TimeoutError:
        proc.kill()
        proc.wait(timeout=2)
        pytest.fail("broker never started accepting connections")
    return proc


def _teardown(proc: subprocess.Popen[bytes]) -> None:
    proc.terminate()
    try:
        proc.wait(timeout=2)
    except subprocess.TimeoutExpired:
        proc.kill()
    if proc.stdout is not None:
        proc.stdout.close()
    if Path(_SOCK).exists():
        Path(_SOCK).unlink()


def _connect_hello(timeout: float = 2.0) -> socket.socket:
    s = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
    s.settimeout(timeout)
    s.connect(_SOCK)
    buf = b""
    while b"\n" not in buf:
        buf += s.recv(4096)
    assert json.loads(buf.split(b"\n", 1)[0]).get("hello") == "mgf-hrb"
    return s


def _rt(s: socket.socket, req: dict) -> dict:
    s.send((json.dumps(req) + "\n").encode())
    buf = b""
    while b"\n" not in buf:
        chunk = s.recv(65536)
        if not chunk:
            break
        buf += chunk
    return json.loads(buf.split(b"\n", 1)[0])


def test_flood_is_shed_with_overloaded() -> None:
    """With a burst of 5 and refill disabled, exactly the first 5 ops
    succeed; every op after that is shed with Err.Overloaded — the
    tenant can't burn unbounded broker work."""
    proc = _spawn({"MGF_HRB_RL_BURST": "5", "MGF_HRB_RL_REFILL_PER_SEC": "0"})
    try:
        s = _connect_hello()
        try:
            outcomes = [_rt(s, {"op": "list", "scope": "hw:/**"}).get("ok") for _ in range(12)]
        finally:
            s.close()
        ok_count = sum(1 for o in outcomes if o is True)
        assert ok_count == 5, (
            f"burst=5, refill=0 should admit exactly 5 ops; got {ok_count} (outcomes={outcomes})"
        )
        # The rest are sheds (ok is False). Confirm the err code.
        assert outcomes[5] is False
    finally:
        _teardown(proc)


def test_overloaded_err_code_and_first_op_unthrottled() -> None:
    """The shed carries the typed `Overloaded` code, and the very first
    op of a fresh tenant is never throttled (the bucket starts full)."""
    proc = _spawn({"MGF_HRB_RL_BURST": "1", "MGF_HRB_RL_REFILL_PER_SEC": "0"})
    try:
        s = _connect_hello()
        try:
            first = _rt(s, {"op": "list", "scope": "hw:/**"})
            second = _rt(s, {"op": "list", "scope": "hw:/**"})
        finally:
            s.close()
        assert first.get("ok") is True, f"first op must not be throttled; got {first!r}"
        assert second.get("ok") is False
        assert second.get("err") == "Overloaded", (
            f"over-rate op must shed with Err.Overloaded; got {second!r}"
        )
    finally:
        _teardown(proc)


def test_bucket_refills_over_time() -> None:
    """Rate-limiting is transient back-pressure: a throttled tenant
    recovers as the bucket refills. burst=2, refill=20/s (1 token per
    50 ms) — exhaust the burst, observe a shed, then after a refill
    window a fresh op succeeds again."""
    proc = _spawn({"MGF_HRB_RL_BURST": "2", "MGF_HRB_RL_REFILL_PER_SEC": "20"})
    try:
        s = _connect_hello()
        try:
            # Drain the burst until we hit a shed (rapid, so refill is
            # negligible during the loop).
            saw_shed = False
            for _ in range(10):
                if _rt(s, {"op": "list", "scope": "hw:/**"}).get("ok") is not True:
                    saw_shed = True
                    break
            assert saw_shed, "a rapid burst past the ceiling must be shed"

            # Wait for the bucket to refill, then a fresh op succeeds.
            time.sleep(0.3)  # ~6 tokens at 20/s, capped at burst=2
            recovered = _rt(s, {"op": "list", "scope": "hw:/**"})
            assert recovered.get("ok") is True, (
                f"the tenant must recover after the bucket refills; got {recovered!r}"
            )
        finally:
            s.close()
    finally:
        _teardown(proc)
