"""Guards for pd_sim.py: the stall shows on the shared layout and the split removes it, the hand-over costs bytes over link speed, the lesson's crossover reproduces,
and the copied toy values still equal lessons 22 and 23."""
import importlib.util
import sys
from pathlib import Path

import pytest

import pd_sim as sim

LESSON_23 = Path(__file__).resolve().parents[2] / "why-ai-waits-before-the-first-word" / "code" / "prefill_sim.py"


def lesson_23():
    spec = importlib.util.spec_from_file_location("prefill_sim_23", LESSON_23)
    module = importlib.util.module_from_spec(spec)
    sys.modules[spec.name] = module            # a dataclass in lesson 23's file looks its module up here
    spec.loader.exec_module(module)
    return module


def test_toy_values_still_equal_lessons_22_and_23():
    ps = lesson_23()
    assert (sim.PARAMS, sim.A100_PEAK_FLOPS, sim.A100_BANDWIDTH) == (ps.PARAMS, ps.A100_PEAK_FLOPS, ps.A100_BANDWIDTH)
    assert sim.CACHE_BYTES_PER_TOKEN == ps.KV_BYTES_PER_TOKEN == 131_072
    assert sim.READ_S_PER_TOKEN == pytest.approx(ps.prefill(1000).math_s / 1000) and sim.STEP_S == pytest.approx(ps.decode_step(1).memory_s)


def test_a_long_prompt_stalls_every_reply_on_a_shared_gpu():
    # Reply A (100-token prompt, 10 words) is mid-reply when B (2,000-token prompt) arrives: A's next word waits for B's whole read plus one step.
    jobs = [(0.0, 100 * sim.READ_S_PER_TOKEN, 9, 0), (0.03, 2000 * sim.READ_S_PER_TOKEN, 3, 1)]
    first, gaps = sim.serve(jobs)
    assert first[0] == pytest.approx(100 * sim.READ_S_PER_TOKEN) and max(gaps[0]) == pytest.approx(2000 * sim.READ_S_PER_TOKEN + sim.STEP_S)
    assert min(gaps[0]) == pytest.approx(sim.STEP_S)


def test_the_hand_over_is_cache_bytes_over_link_speed():
    assert sim.handover_seconds(2000, 25) == 2000 * 131_072 / 25e9
    one = [sim.Request(1.0, 2000, 4)]
    first, gaps = sim.split(one, 1.25)
    assert first[0] - 1.0 == pytest.approx(2000 * sim.READ_S_PER_TOKEN + sim.handover_seconds(2000, 1.25))      # first word reaches the reader with the cache
    assert gaps[0] == pytest.approx([sim.STEP_S] * 3)


def test_split_removes_the_stall_and_the_lessons_crossover_reproduces():
    requests = sim.make_requests(4, 2000, 200)
    shared = sim.measure(requests, sim.shared(requests))
    fast = sim.measure(requests, sim.split(requests, 25))
    assert shared["gap_p99"] > 0.1 > 10 * fast["gap_p99"] and fast["gap_p99"] <= 2 * sim.STEP_S and shared["gap_p50"] < 0.01      # p99 stalls, p50 does not
    assert fast["goodput"] > 0.9 > 0.2 > shared["goodput"] and fast["first_p99"] > shared["first_p99"]                           # the split wins goodput, loses first word
    slow = sim.measure(requests, sim.split(requests, 0.125))
    assert slow["goodput"] < shared["goodput"]                                                                                   # a slow enough link ends the win
    short = sim.make_requests(4, 200, 200)
    s, f = sim.measure(short, sim.shared(short)), sim.measure(short, sim.split(short, 25))
    assert s["goodput"] == f["goodput"] == 1.0 and f["first_p99"] >= s["first_p99"] * 0.7                                        # short prompts: a tie