"""Guards for tail_sim.py: the toy server against lesson 24's, the percentile rule, the session chance, the two testers, and that the numbers the lesson
and its figures print are the ones the code gives."""
import ast
import re
from functools import lru_cache
from pathlib import Path

import numpy as np
import pytest

import make_figs
import tail_sim as sim

HERE = Path(__file__).resolve().parent
LESSON = HERE.parent
W = sim.WORKLOAD


@lru_cache(maxsize=None)
def run(rate, connections=0, pause=0.0):
    return sim.serve(W, rate, connections, pause)


@lru_cache(maxsize=None)
def swept():
    return sim.sweep()


def test_toy_server_gives_the_times_of_lesson_24_continuous_batching():
    """Guards the toy drifting from the scheduler lesson 24 explained: on schedule it must match that lesson's `simulate` for every request."""
    b = sim.batching
    for rate in (4, 12, 18):
        intended = sim.intended_times(W["gaps"], rate)
        reference = b.simulate([b.Request(t, p, r) for t, p, r in zip(intended, W["prompts"], W["replies"])], sim.SLOTS, "continuous")
        mine = run(rate)
        assert mine.first == reference.first_token and mine.done == reference.finished and mine.steps == len(reference.steps), rate


def test_nearest_rank_percentile_and_how_many_requests_a_p99_rests_on():
    """Guards an off-by-one in the percentile: p99 of 2,000 is the 1,980th time (20 slower), of 100 the 99th (1 slower), of 50 the slowest."""
    assert [sim.percentile(list(range(1, 101)), p) for p in (50, 95, 99)] == [50, 95, 99]
    assert sim.percentile(list(range(2000, 0, -1)), 99) == 1980                   # unsorted input
    assert sum(1 for v in range(1, 2001) if v > sim.percentile(list(range(1, 2001)), 99)) == 20
    assert sim.percentile(list(range(1, 51)), 99) == 50 and sim.percentile([7.0], 99) == 7.0
    assert sim.percentile([1, 2, 3, 4], 50) == 2                                  # nearest rank takes the lower of the two middle values
    times = sim.totals(run(sim.OPENING_RATE))                                     # NumPy's inverted_cdf is nearest rank; its default interpolates, and differs by under a thousandth here
    for p in sim.PERCENTILES:
        assert np.percentile(times, p, method="inverted_cdf") == sim.percentile(times, p)
        assert 0 < abs(np.percentile(times, p) - sim.percentile(times, p)) < 1e-3
    assert round(float(np.percentile(range(1, 51), 99)), 2) == 49.51 and sim.percentile(list(range(1, 51)), 99) == 50


def test_the_average_sits_between_the_median_and_p95_and_p99_grows_with_load():
    """Guards the lesson's central claim and the order of the numbers: p50 < average < p95 < p99 < max for both measures at every load, and the load stops stay below capacity."""
    assert sim.RATES[-1] < sim.capacity() < sim.RATES[-1] * 1.05
    for measure in ("total", "wait"):
        rows = [swept()[rate][measure] for rate in sim.RATES]
        for row in rows:
            assert row["p50"] < row["average"] < row["p95"] < row["p99"] < row["max"], (measure, row)
        for before, after in zip(rows, rows[1:]):
            assert after["p99"] > before["p99"] and after["average"] > before["average"]


def test_session_chance_is_one_minus_the_chance_of_missing_every_slow_request():
    """Guards the sum 10 x 5% = 50% (wrong) in place of 1 - 0.95^10 = 40.1%; the Tail at Scale's 63% for 100 servers at 1 in 100 is the same arithmetic."""
    assert round(sim.session_chance(0.05, 10), 4) == 0.4013 and round(sim.session_chance(0.05, 20), 4) == 0.6415
    assert round(sim.session_chance(0.01, 100), 2) == 0.63
    product = 1.0
    for _ in range(20):
        product *= 0.95
    assert sim.session_chance(0.05, 20) == pytest.approx(1 - product)
    assert sim.session_chance(0.05, 1) == pytest.approx(0.05) and sim.session_chance(0.0, 50) == 0.0


def test_a_tester_that_waits_for_replies_under_reports_the_stall():
    """Guards the lesson's second warning: with the same 5 s freeze at 12 a second, the tester that waits for each reply sees about 8 requests stuck and a
    p99 wait near the normal one, while the tester on schedule sees about a tenth of all requests wait over a second. Also: neither tester sends early,
    and the waiting one never has more than 8 requests out."""
    demo = sim.generator_demo()
    schedule, waits = demo["on schedule"], demo["waits for each reply"]
    assert schedule["wait"]["p99"] > 100 * waits["wait"]["p99"] and waits["wait"]["p99"] < 0.1 and waits["wait"]["max"] > sim.PAUSE - 0.5
    assert schedule["slow"] >= 100 and waits["slow"] <= 10
    for connections in (0, sim.CONNECTIONS):
        r = run(sim.DEMO_RATE, connections, sim.PAUSE)
        assert all(s >= i for s, i in zip(r.sent, r.intended))
        if connections:
            events = sorted([(s, 1) for s in r.sent] + [(d, -1) for d in r.done], key=lambda e: (e[0], e[1]))      # a reply that completes at time t frees its connection first
            out = peak = 0
            for _, delta in events:
                out += delta
                peak = max(peak, out)
            assert peak <= sim.CONNECTIONS
            assert all(r.sent[h + connections] >= r.done[h] for h in range(sim.N_REQUESTS - connections))


def test_averaging_two_p99s_is_not_the_p99_of_all_the_requests():
    """Guards the lesson's averaged-percentile example: two servers' p99s averaged give about half of what all 4,000 requests together give."""
    ex = sim.pooled_example()
    quiet, busy = sim.waits(run(4)), sim.waits(run(19))
    pooled = sorted(quiet + busy)[-(-99 * 4000 // 100) - 1]
    assert ex["pooled"] == pooled and ex["averaged"] == (ex["quiet_p99"] + ex["busy_p99"]) / 2 and ex["requests"] == 2000
    assert ex["averaged"] < 0.6 * ex["pooled"]
    assert sorted(quiet + busy)[-40:] == sorted(busy)[-40:]                       # the 40 slowest of the 4,000 all come from the busy server


def test_every_shuffle_is_the_same_requests_in_another_order():
    """Guards a shuffle map that repeats or drops a request: the widget's rerolls must be permutations of the 2,000 requests."""
    assert sim.REROLLS[0] == (1, 0) and len(set(sim.REROLLS)) == len(sim.REROLLS)
    for a, b in sim.REROLLS:
        mixed = sim.shuffled(W, a, b)
        assert sorted(mixed["prompts"]) == sorted(W["prompts"]) and sorted(mixed["replies"]) == sorted(W["replies"]) and sorted(mixed["gaps"]) == sorted(W["gaps"])
    assert sim.shuffled(W, 1, 0) == W


# The only modules this folder imports. None of them reaches a network, so a module that is not on the list (socketserver, xmlrpc, a new HTTP library) fails
# the test until someone adds it here on purpose. subprocess is handled apart: tests may import it.
IMPORTS = {"__future__", "ast", "collections", "functools", "heapq", "importlib", "json", "math", "numpy", "pathlib", "pytest", "random", "re", "sys", "typing", "xml"}
CALLS = {"urlopen", "connect", "system", "popen", "__import__", "import_module"}


def test_the_code_asks_nothing_of_the_network():
    """Guards the hard rule of this lesson: a load generator that only exists inside this process. No code file, the tests included, imports a module outside
    IMPORTS or calls one that reaches a network, a shell or an import by name; the one subprocess call allowed is `subprocess.run(["node", ...])` in a test, which
    test_widget.py uses to run the widget's own script. Each failure is a sentence that names the file, what it found and the rule."""
    own, found = {p.stem for p in HERE.glob("*.py")}, []
    for path in sorted(HERE.glob("*.py")):
        is_test = path.name.startswith("test_")
        tree = ast.parse(path.read_text(encoding="utf-8"))
        aliases = {a.asname or a.name for n in ast.walk(tree) if isinstance(n, ast.Import) for a in n.names if a.name == "subprocess"}      # `import subprocess as sp` too
        for node in ast.walk(tree):
            if isinstance(node, (ast.Import, ast.ImportFrom)):
                for module in [a.name for a in node.names] if isinstance(node, ast.Import) else [node.module or ""]:
                    root = module.split(".")[0]
                    if not (root in IMPORTS | own or (root == "subprocess" and is_test and isinstance(node, ast.Import))):
                        found.append(f"{path.name} imports {module or 'a relative module'}, but only the modules in IMPORTS may be imported here, and subprocess only in a test.")
            elif isinstance(node, ast.Call):
                f = node.func
                name = f.attr if isinstance(f, ast.Attribute) else f.id if isinstance(f, ast.Name) else ""
                if isinstance(f, ast.Attribute) and isinstance(f.value, ast.Name) and f.value.id in aliases:
                    program = node.args[0].elts[0] if node.args and isinstance(node.args[0], ast.List) and node.args[0].elts else None
                    if not (is_test and name == "run" and isinstance(program, ast.Constant) and program.value == "node"):
                        found.append(f"{path.name} calls subprocess.{name}(), but only subprocess.run([\"node\", ...]) is allowed, in a test.")
                elif name in CALLS:
                    found.append(f"{path.name} calls {name}(), which reaches a network, a shell or an import by name; only subprocess.run([\"node\", ...]) is allowed, in a test.")
            elif isinstance(node, ast.Constant) and isinstance(node.value, str) and not is_test and re.search(r"https?://", node.value):      # the tests name URLs in their lists of banned words
                found.append(f"{path.name} contains the URL text {node.value[:40]!r}, but only a test may name a URL, in its list of banned words.")
    assert not found, "network or process code in the lesson folder:\n  " + "\n  ".join(found)


def fragments():
    """Every number the lesson text quotes, as the text writes it, computed from the code."""
    f, S = sim.seconds_text, swept()
    total, wait = S[17]["total"], S[17]["wait"]
    times = sim.totals(run(17))
    above = sum(1 for v in times if v > total["average"])
    ex, demo, spread = sim.pooled_example(), sim.generator_demo(), sim.shuffle_spread()
    schedule, waits = demo["on schedule"], demo["waits for each reply"]
    waiting = run(12, sim.CONNECTIONS, sim.PAUSE)
    late = lambda r: sum(1 for s, i in zip(r.sent, r.intended) if s > i)
    froze = lambda r: sum(1 for s in r.sent if sim.PAUSE_AT <= s < sim.PAUSE_AT + sim.PAUSE)
    shoppers = [1] * 8 + [16] * 2
    out = [
        f"takes {f(total['average'])} seconds (s) on average, yet 1 request in 20 takes more than {f(total['p95'])} s and 1 in 100 more than {f(total['p99'])} s",
        f"Eight wait 1 minute and two wait 16, so the average is {sim.mean(shoppers):g} minutes", f"the 5th of 10 sorted, which is {sim.percentile(shoppers, 50)} minute",
        f"near {sim.capacity():.1f}", f"the total time averages {f(total['average'])} s, but p95 is {f(total['p95'])} s and p99 is {f(total['p99'])} s",
        f"{len(times) - above:,} of the {len(times):,} requests take the average of {f(total['average'])} s or less and {above:,} take longer",
        f"{sum(1 for r in W['replies'] if r >= sim.LONG_REPLY[0])} of the {len(times):,} are long, {sim.LONG_REPLY[0]} tokens or more, against at most {sim.SHORT_REPLY[1]}",
        f"p50 {f(total['p50'])} s", f"slowest {f(total['max'])} s",
        f"{100 * sim.session_chance(0.05, 10):.0f}%", f"{100 * sim.session_chance(0.05, 20):.0f}%",
        f"{100 * sim.session_chance(0.05, 10):.1f}% for 10 and {100 * sim.session_chance(0.05, 20):.1f}% for 20", f"which is {100 * sim.session_chance(0.01, 100):.1f}%",
        f"and the p99 wait is {f(wait['p99'])} s, against {f(total['average'])} s and {f(total['p99'])} s for the total",
        f"average wait is {f(wait['average'])} s",
        f"p99 wait is {f(S[4]['wait']['p99'])} s at 4 a second and {f(S[19]['wait']['p99'])} s at 19",
        f"At 12 requests a second the average is {f(S[12]['wait']['average'])} s, p95 is {f(S[12]['wait']['p95'])} s and p99 is {f(S[12]['wait']['p99'])} s",
        f"At 19 the average is {f(S[19]['wait']['average'])} s and p99 is {f(S[19]['wait']['p99'])} s",
        f"from {f(S[4]['wait']['p50'])} to {f(S[12]['wait']['p50'])} s, while the p99 rises from {f(S[4]['wait']['p99'])} to {f(S[12]['wait']['p99'])} s",
        f"an average of {f(S[12]['wait']['average'])} s looks healthy while 1 request in 20 already waits more than {f(S[12]['wait']['p95'])} s",
        f"the p99 wait is {f(schedule['wait']['p99'])} s for a tester on schedule and {f(waits['wait']['p99'])} s for one that waits",
        f"{f(S[4]['wait']['p99'])} s and {f(S[19]['wait']['p99'])} s average to {f(ex['averaged'])} s, but all {2 * ex['requests']:,} requests together have a p99 of {f(ex['pooled'])} s",
        f"passes at {f(total['average'])} s while 1 request in 20 takes over {f(total['p95'])} s",
        f"Their p99 waits are {ex['quiet_p99']:.4f} s and {ex['busy_p99']:.4f} s, and the average of the two is {ex['averaged']:.4f} s", f"the p99 is {ex['pooled']:.4f} s",
        "The 40 slowest of the 4,000 all come from the busy server", f"a p99 of {f(S[4]['total']['p99'])} s",
        f"over its {len(sim.REROLLS)} orders at 17 a second the average total time runs from {f(spread['average'][0])} to {f(spread['average'][1])} s "
        f"and the p99 from {f(spread['p99'][0])} to {f(spread['p99'][1])} s",
        f"seed {sim.SEED}", f"prompts of {min(W['prompts'])} to {max(W['prompts'])} tokens and replies of {min(W['replies'])} to {max(W['replies'])}",
        f"| On schedule | {froze(run(12, 0, sim.PAUSE))} | {f(schedule['wait']['p99'])} s |", f"| Waiting | {froze(waiting)} | {f(waits['wait']['p99'])} s |",
        f"{waits['slow']} of its requests wait over {sim.SLOW_WAIT:g} s for the first word, against {schedule['slow']} for the tester on schedule",
        f"Only {froze(waiting)} requests go out in the freeze instead of {froze(run(12, 0, sim.PAUSE))}. Even with no freeze, {late(run(12, sim.CONNECTIONS)):,} of its "
        f"{sim.N_REQUESTS:,} requests go out later than planned, because the last reply on their connection is not yet complete at the planned time; "
        f"the freeze adds {late(waiting) - late(run(12, sim.CONNECTIONS)):,}, so {late(waiting):,} in all",
        f"Its one worst wait, {f(waits['wait']['max'])} s", f"the server freezes for {sim.PAUSE:g} s at {sim.PAUSE_AT:g} s", f"at {sim.DEMO_RATE} requests a second",
        f"here {sim.CONNECTIONS}, and each sends",
    ]
    for rate in sim.RATES:
        for measure in ("total", "wait"):
            s = S[rate][measure]
            out.append(f"| {rate} | " + " | ".join(f(s[k]) for k in ("average", "p99")) + " |")
    for n in (50, 100, 200, 1000, 2000):
        rank = sim.percentile(list(range(1, n + 1)), 99)
        out.append(f"| {n:,} | {rank:,} | {n - rank:,} |")
    return out


def test_the_lesson_and_its_figures_quote_the_numbers_the_code_gives():
    """Guards a stale page: when the seed, a constant or the server changes, every number the prose, tables and figures quote must change with it. One assert
    names every fragment that is missing, so a rerun does not have to be repeated for each. Also: the figures are the ones make_figs.py draws, and only the
    site's colours are in them (no blue, no vermillion)."""
    text = (LESSON / "lesson.md").read_text(encoding="utf-8")
    missing = [frag for frag in fragments() if frag not in text]
    for name, draw in make_figs.FIGURES.items():
        path = make_figs.FIGS / name
        if not path.exists() or path.read_text(encoding="utf-8") != draw():
            missing.append(f"figs/{name} differs from make_figs.py")
        elif f"figs/{name}" not in text:
            missing.append(f"lesson.md does not show figs/{name}")
    assert not missing, "missing from the lesson or its figures:\n  " + "\n  ".join(missing)
    allowed = {"#0a1626", "#44566b", "#b36200", "#00855f", "#d2dae4", "#ebf0f6", "#ffffff"}
    for path in sorted((LESSON / "figs").glob("*.svg")):
        colours = {c.lower() for c in re.findall(r"#[0-9a-fA-F]{6}\b", path.read_text(encoding="utf-8"))}
        assert colours <= allowed, (path.name, colours - allowed)