"""Tests that tie lesson.md, trace.py and the four figures to pd_sim.py. Run from this folder: python3 -B -m pytest -q -p no:cacheprovider"""
import io
import re
import subprocess
import sys
import xml.etree.ElementTree as ET
from contextlib import redirect_stdout
from functools import lru_cache
from pathlib import Path

import make_figs
import pd_sim as sim

HERE = Path(__file__).resolve().parent
LESSON = HERE.parent
SOURCE = (HERE / "pd_sim.py").read_text(encoding="utf-8")
RPS, PROMPT, REPLY, LINK = make_figs.RPS, make_figs.PROMPT, make_figs.REPLY, make_figs.LINK
ms = lambda seconds: f"{seconds * 1000:.1f}"
pct = lambda share: f"{share:.1%}"
READ = PROMPT * sim.READ_S_PER_TOKEN


@lru_cache(maxsize=None)
def reading():
    prose = (LESSON / "lesson.md").read_text(encoding="utf-8")
    return prose, {name: (LESSON / "figs" / name).read_text(encoding="utf-8") for name in make_figs.FIGURES}


def test_every_number_the_lesson_and_its_figures_quote_comes_from_the_code():
    """Guards prose, trace and figure numbers drifting from the code: each fragment is built from a value computed here and must appear in the lesson or the figure.
    The numbers quoted from the papers and the pull request are not computed here."""
    prose, figures = reading()
    run = make_figs.runs(RPS, PROMPT, REPLY, LINK)
    both, apart = run["shared"], run["split"]
    at = lambda rps=RPS, prompt=PROMPT, link=LINK: make_figs.runs(rps, prompt, REPLY, link)
    requests = sim.make_requests(RPS, PROMPT, REPLY)
    _, gaps = sim.shared(requests)
    stalled = sum(1 for i in range(len(requests)) if max(gaps[i], default=0.0) > sim.WORD_GAP_TARGET_S) / len(requests)
    load = sum(r.prompt for r in requests) * sim.READ_S_PER_TOKEN / requests[-1].arrival
    cache = PROMPT * sim.CACHE_BYTES_PER_TOKEN
    break_even = sim.CACHE_BYTES_PER_TOKEN / sim.READ_S_PER_TOKEN / 1e9
    stall = READ + sim.STEP_S
    cross = make_figs.crossover_gbps()
    shortest = round(PROMPT * min(sim.MIX))
    over = next(rps for rps in make_figs.RPS_GRID if at(rps)["shared"]["gap_p99"] > 0.1)
    slow = at(link=0.25)["split"]
    scene = make_figs.stall_scene()
    wanted = {"lesson.md": [
        f"a prompt of {PROMPT} tokens", f"takes {ms(READ)} ms to read", f"a word due {ms(sim.STEP_S)} ms after the last comes {ms(stall)} ms after it",
        f"At {RPS} requests a second only {pct(both['goodput'])} of requests", f"within {sim.FIRST_WORD_TARGET_S} s and no later gap over {sim.WORD_GAP_TARGET_S} s",
        f"and {pct(apart['goodput'])[:-1]}% do", f"The price is that the reading GPU's notes on the prompt (the cache, {cache / 1e6:.0f} MB for {PROMPT} tokens) must cross the link between the two GPUs",
        f"{RPS} requests a second, prompts of {PROMPT} tokens, replies of {REPLY} words and a {LINK} GB/s link", f"shared GPUs get {pct(both['goodput'])} of requests right and the split {pct(apart['goodput'])}",
        f"Drag Requests a second to 10: shared falls to {pct(at(10)['shared']['goodput'])} and the split to {pct(at(10)['split']['goodput'])}",
        f"Drag Link speed to 0.25 GB/s: the split falls to {pct(slow['goodput'])}", f"Drag Prompt tokens to 200: the two tie at {pct(at(prompt=200)['split']['goodput'])}",
        f"serves {sim.N_REQUESTS} requests", f"600 requests come at {RPS} a second",
        f"prompts of {round(PROMPT * min(sim.MIX))} to {round(PROMPT * max(sim.MIX))} tokens and replies of {round(REPLY * min(sim.MIX))} to {round(REPLY * max(sim.MIX))} words",
        f"Half of all gaps between words are {ms(both['gap_p50'])} ms (the p50), but the slowest 1 in 100 (the p99, the 99th percentile) is {ms(both['gap_p99'])} ms",
        f"Only {pct(both['goodput'])} of requests are good", f"with {load:.2f} s of reading to do each second",
        f"{PROMPT} x {sim.CACHE_BYTES_PER_TOKEN:,} bytes = {cache / 1e6:.1f} MB", f"At {LINK} GB/s that takes {ms(sim.handover_seconds(PROMPT, LINK))} ms, against {ms(READ)} ms to read",
        f"no gap passes {ms(apart['gap_p99'])} ms and {pct(apart['goodput'])} of requests are good, against {pct(both['goodput'])}",
        f"p50 {ms(apart['first_p50'])} ms against {ms(both['first_p50'])} ms", f"The word waits {ms(stall)} ms, {round(stall / sim.STEP_S)} times the usual {ms(sim.STEP_S)}",
        f"A 200-word reply lasts about {REPLY * sim.STEP_S:.1f} s", f"so {pct(stalled)} of replies have a gap over 0.1 s and {pct(both['goodput'])} are good",
        f"passes 100 ms at {over} requests a second and is {ms(at(10)['shared']['gap_p99'])} ms at 10", f"One reader can read about {1 / READ:.1f} prompts a second (1 / {READ:.3f} s)",
        f"from {at(1)['split']['first_p99']:.2f} s at 1 request a second to {at(10)['split']['first_p99']:.2f} s at 10", f"equals the read at {break_even:.2f} GB/s for any prompt length",
        f"p99 {ms(apart['first_p99'])} ms against {ms(both['first_p99'])} ms", f"{pct(at(link=1.25)['split']['goodput'])} at 1.25 GB/s and {pct(at(link=0.125)['split']['goodput'])} at 0.125",
        f"on links faster than about {cross:.1f} GB/s", f"the shortest prompt, {shortest} tokens, has a {shortest * sim.CACHE_BYTES_PER_TOKEN / 1e6:.0f} MB cache that cannot be read and cross within {sim.FIRST_WORD_TARGET_S} s",
        "change line 8 to `LINK_GBPS = 0.25`", f"The hand-over takes {ms(sim.handover_seconds(PROMPT, 0.25))} ms, about ten times the read",
        f"Even the shortest prompt, {shortest} tokens, needs {sim.handover_seconds(shortest, 0.25) * 1000:.0f} ms to cross",
        f"The slowest 1 in 100 is {ms(both['gap_p99'])} ms shared and {ms(apart['gap_p99'])} ms split", f"at {PROMPT} tokens the split leads {pct(apart['goodput'])} to {pct(both['goodput'])}",
        f"the split stops paying below about {cross:.1f} GB/s", f"(0.2922 in the code: at 0.29 no request is good and at 0.293, {pct(at(link=0.293)['split']['goodput'])} are)",
        f"{ms(shortest * sim.READ_S_PER_TOKEN)} ms to read plus {ms(sim.handover_seconds(shortest, cross))} ms to cross fills the {sim.FIRST_WORD_TARGET_S} s",
        f"`pd_sim.py` is {len(SOURCE.splitlines())} lines", f"{sim.READ_S_PER_TOKEN * 1e6:.1f} microseconds to read a prompt token and {sim.STEP_S * 1000:.3f} ms for a writing step",
        f"The {sim.CACHE_BYTES_PER_TOKEN:,} cache bytes per token", f"the {sim.SLOTS} slots", f"For a {PROMPT}-token prompt, {cache / 1e6:.1f} MB crosses",
        f"at 2 GB/s the split is good for {pct(at(link=2)['split']['goodput'])} of requests", f"at 200 tokens the read takes {ms(200 * sim.READ_S_PER_TOKEN)} ms, a stall is about {round((200 * sim.READ_S_PER_TOKEN + sim.STEP_S) * 1000)} ms",
        f"6000-token prompts need {6000 * sim.READ_S_PER_TOKEN:.2f} s of reading each", f"(the split's first-word p99 is {at(prompt=6000)['split']['first_p99']:.1f} s)",
        f"takes {sim.handover_seconds(PROMPT, 2) * 1000:.0f} ms, longer than the {READ * 1000:.0f} ms read", f"{pct(at(link=2)['split']['goodput'])} of requests at 2 GB/s, against {pct(both['goodput'])} shared",
        f"{pct(stalled)} of the shared requests have one", f"A 200-token prompt takes {ms(200 * sim.READ_S_PER_TOKEN)} ms to read", f"about {round((200 * sim.READ_S_PER_TOKEN + sim.STEP_S) * 1000)} ms",
        f"first-word p99 reaches {at(10)['split']['first_p99']:.2f} s", f"An average {PROMPT}-token prompt takes {READ:.3f} s to read, so one reader handles about {1 / READ:.1f} prompts a second",
    ], "stall.svg": [
        f"{ms(stall)} ms with no word from A", f"reads B: {ms(READ)} ms",
        f"B's cache crosses: {ms(sim.handover_seconds(PROMPT, LINK))} ms", f"A: a word every {ms(sim.STEP_S)} ms",
    ], "gap-by-load.svg": [f"{ms(at(rps)['shared']['gap_p99'])} at {rps}" for rps in (1, 3, 4, 6, 10)] + [f"{ms(at(10)['split']['gap_p99'])} at 10"],
        "goodput-by-link.svg": [f"{pct(at(link=link)['split']['goodput'])} at {link}" for link in sim.LINKS_GBPS] + [pct(both["goodput"]), f"crossover at about {cross:.1f} GB/s", f"about {cross:.1f} GB/s</text>"],
        "goodput-by-prompt.svg": [f"{pct(at(prompt=p)['shared']['goodput'])} at {p:,}" for p in (200, 1000, 2000, 4000)] + [f"{pct(at(prompt=p)['split']['goodput'])} at {p:,}" for p in (200, 1000, 2000, 4000, 6000)],
    }
    texts = {"lesson.md": prose, **figures}
    missing = [f"{name}: {fragment!r}" for name, fragments in wanted.items() for fragment in fragments if fragment not in texts[name]]
    assert not missing, "not found: " + "; ".join(missing)
    assert len(SOURCE.splitlines()) < 150 and f"{cross:.4f}" == "0.2922" and at(link=0.29)["split"]["goodput"] == 0.0 and stalled == 1 - both["goodput"]


def test_the_lessons_claims_about_where_the_split_pays_hold():
    """Guards the shape claims: the layouts tie up to 600 tokens and the split leads from 800 to 4000 and both collapse at 6000; a link under the crossover leaves no request
    able to meet the first-word target even for the shortest prompt, and one just over it wins; the split loses the first word at every link speed on offer."""
    for prompt in (100, 200, 400, 600):
        c = make_figs.runs(RPS, prompt, REPLY, LINK)
        assert abs(c["split"]["goodput"] - c["shared"]["goodput"]) < 0.01, prompt
    for prompt in (800, 1000, 1500, 2000, 3000, 4000):
        c = make_figs.runs(RPS, prompt, REPLY, LINK)
        assert c["split"]["goodput"] - c["shared"]["goodput"] > 0.02, prompt
    both = make_figs.runs(RPS, 6000, REPLY, LINK)
    assert max(both["shared"]["goodput"], both["split"]["goodput"]) < 0.05
    cross, shortest = make_figs.crossover_gbps(), round(PROMPT * min(sim.MIX))
    assert abs(shortest * sim.READ_S_PER_TOKEN + sim.handover_seconds(shortest, cross) - sim.FIRST_WORD_TARGET_S) < 0.002
    assert make_figs.runs(RPS, PROMPT, REPLY, cross * 0.97)["split"]["goodput"] == 0.0
    assert make_figs.runs(RPS, PROMPT, REPLY, cross * 1.03)["split"]["goodput"] > make_figs.runs(RPS, PROMPT, REPLY, LINK)["shared"]["goodput"]
    assert all(make_figs.runs(RPS, PROMPT, REPLY, link)["split"]["first_p50"] > make_figs.runs(RPS, PROMPT, REPLY, link)["shared"]["first_p50"] for link in sim.LINKS_GBPS)


def test_the_figures_are_the_ones_the_code_draws_and_follow_the_house_style():
    """Guards a figure going stale or off-palette: the files on disk equal what make_figs draws, the lesson shows each one, and each parses, has a title and a description,
    a white ground, no script, no web address, only the house colours and no text under 20 units."""
    prose, figures = reading()
    assert set(figures) == set(re.findall(r"\]\(figs/([\w-]+\.svg)\)", prose)) == set(make_figs.FIGURES)
    allowed = {make_figs.INK, make_figs.GREY, make_figs.ORANGE, make_figs.GREEN, make_figs.FILL, make_figs.PANEL, "#ffffff"}
    web_address = re.compile("ht" + r"tps?://")
    for name, svg in figures.items():
        assert svg == make_figs.FIGURES[name](), name
        root = ET.fromstring(svg)
        assert root.find("{%s}title" % make_figs.NAMESPACE) is not None and root.find("{%s}desc" % make_figs.NAMESPACE) is not None and 'fill="#ffffff"' in svg, name
        assert "<script" not in svg and not web_address.search(svg.replace(make_figs.NAMESPACE, "")) and root.get("width") == "540", name
        assert {c.lower() for c in re.findall(r"#[0-9a-fA-F]{6}\b", svg)} <= {c.lower() for c in allowed}, name
        assert all(int(size) >= 20 for size in re.findall(r'font-size="(\d+)"', svg)), name
    for path in HERE.glob("*.py"):
        assert not web_address.search(path.read_text(encoding="utf-8").replace("ht" + "tp://www.w3.org", "")), path.name


def table_rows(prose):
    return re.findall(r"^\| \d\. [^|]+\| `(\w+\.py):(\d+)`<br>`(.+?)` \|", prose, re.M)


def test_code_blocks_are_slices_of_the_real_file_and_the_table_points_at_its_lines():
    """Guards a `file:line` in the code table pointing at the wrong line, and a code block in Go deeper that no longer matches the file."""
    prose, _ = reading()
    blocks = re.findall(r"```python\n(.*?)\n```", prose, re.S)
    assert len(blocks) == 4 and all(block in SOURCE for block in blocks)
    rows = table_rows(prose)
    assert [(name, int(n)) for name, n, _ in rows] == [("pd_sim.py", 41), ("pd_sim.py", 57), ("pd_sim.py", 72), ("pd_sim.py", 98), ("pd_sim.py", 80), ("pd_sim.py", 112)]
    for name, number, code in rows:
        line = (HERE / name).read_text(encoding="utf-8").splitlines()[int(number) - 1]
        assert re.sub(r"\s+# .*$", "", line).strip() == code, (name, number)


def run_trace(source):
    out = io.StringIO()
    with redirect_stdout(out):
        exec(compile(source, "trace.py", "exec"), {"__name__": "__main__"})
    return out.getvalue()


def test_trace_is_the_block_in_the_lesson_and_the_your_turn_answer_is_right():
    """Guards the pasted run drifting from `python3 trace.py` (the real command, from the code folder) and the Your turn answer: changing line 8 to LINK_GBPS = 0.25 must
    give the lines in the details block."""
    prose, _ = reading()
    real = subprocess.run([sys.executable, "-B", "trace.py"], cwd=HERE, capture_output=True, text=True, check=True).stdout
    pasted = re.search(r"python3 trace\.py\n```\n\nPhone readers: slide the box sideways to see every column\.\n\n```text\n(.*?)\n```", prose, re.S).group(1)
    assert real == pasted + "\n"
    trace = (HERE / "trace.py").read_text(encoding="utf-8")
    assert trace.splitlines()[7].startswith("LINK_GBPS = 25 ") and "change line 8 to `LINK_GBPS = 0.25`" in prose
    changed = run_trace(trace.replace("LINK_GBPS = 25 ", "LINK_GBPS = 0.25 ", 1))
    answer = re.search(r"<summary>What you should see</summary>.*?```text\n(.*?)\n```", prose, re.S).group(1)
    assert answer + "\n" == changed and "goodput 0.0% against 13.5% shared" in answer
    assert [a == b for a, b in zip(real.splitlines(), changed.splitlines())] == [True, True, True, True, False, False, True]