"""The panels. Every number printed here is measured, in this process.

    python3 demo.py

Panel 6 spawns OS threads until the operating system refuses, which takes
about twenty seconds. Everything before it is instant.
"""
import asyncio
import gc
import os
import platform
import resource
import sys
import threading
import time
import tracemalloc
import traceback
import warnings

import runtimes as rt
from runtimes import DEPOSITS, START, TELLERS, run_chain, run_green, run_stackful

SEP = "=" * 70
EXPECT = START + TELLERS * DEPOSITS


def banner():
    print("stackless-vs-stackful: the suspension point, and what it costs")
    print(SEP)
    print(" ", platform.platform(), platform.machine())
    print("  Python", sys.version.replace("\n", " "))
    gil = getattr(sys, "_is_gil_enabled", None)
    print(f"  GIL: {'enabled' if gil and gil() else 'n/a'}"
          f"   switchinterval={sys.getswitchinterval() * 1e3:.1f} ms")
    print()
    print(f"  workload: {TELLERS} tellers x {DEPOSITS} deposits of 1 into one")
    print(f"  account starting at {START}. The only correct answer is {EXPECT}.")
    print("  The critical section is five frames deep.")
    print()


# ---------------------------------------------------------------- 1


def panel_headline():
    print("1. THE HEADLINE -- one line is added to a library function")
    print(SEP)
    print("  ACT 1 -- library v1: io_read does not suspend")
    bal, sw = run_stackful("none")
    print(f"    stackful  (OS threads)        final={bal:<6} "
          f"switches={sw:<5} {'OK' if bal == EXPECT else 'WRONG'}")
    bal, sw = run_green(rt.g1_teller)
    print(f"    stackless (generators)        final={bal:<6} "
          f"switches={sw:<5} {'OK' if bal == EXPECT else 'WRONG'}")
    print()
    print("  ACT 2 -- library v2: io_read gains one line and starts suspending")
    print("           between reading the row and returning it. No frame above")
    print("           it is edited, in either column.")
    bal, sw = run_stackful("after_read")
    print(f"    stackful  (OS threads)        final={bal:<6} "
          f"switches={sw:<5} {'OK' if bal == EXPECT else 'WRONG'}")
    try:
        bal, sw = run_green(rt.g2_teller)
        print(f"    stackless (generators)        final={bal:<6} "
              f"switches={sw:<5} {'OK' if bal == EXPECT else 'WRONG'}")
    except Exception as e:
        frame = traceback.extract_tb(e.__traceback__)[-1]
        print("    stackless (generators)        REFUSES TO RUN")
        print(f"        {os.path.basename(frame.filename)}, "
              f"line {frame.lineno}, in {frame.name}")
        print(f"          {frame.line}")
        print(f"        {type(e).__name__}: {e}")
    print()
    print("  ACT 3 -- stackless, after colouring every frame above io_read")
    bal, sw = run_green(rt.g3_teller)
    print(f"    stackless (generators, fixed) final={bal:<6} "
          f"switches={sw:<5} {'OK' if bal == EXPECT else 'WRONG'}")
    print()
    print("    frames that had to be edited to make the program run again:")
    print("      stackful:  0")
    print("      stackless: 4  (db_fetch, repo_get, deposit, teller)")
    print()
    print("    Colouring bought visibility. It did not buy correctness.")
    print()


# ---------------------------------------------------------------- 2


def bank(where, tellers, deposits):
    return run_stackful(where, tellers, deposits)[0]


def panel_arithmetic():
    print(f"2. WHERE {EXPECT - 50} COMES FROM")
    print(SEP)
    print("  Round-robin lockstep: in each round every teller reads the same")
    print("  balance v and writes v+1, so a ROUND advances the balance by 1")
    print("  no matter how many tellers are in it.")
    print(f"    {DEPOSITS} rounds  =>  {START} + {DEPOSITS} = {START + DEPOSITS}")
    print(f"    {TELLERS * DEPOSITS} increments were performed, "
          f"{TELLERS * DEPOSITS - DEPOSITS} of them were lost.")
    print()
    print("  predicted final = START + deposits, independent of teller count:")
    print("    tellers  deposits  increments  predicted  observed")
    for t in (2, 3, 4):
        for d in (10, 50):
            pred, got = START + d, bank("after_read", t, d)
            flag = "" if pred == got else "   <-- MISMATCH"
            print(f"    {t:<8} {d:<9} {t * d:<11} {pred:<10} {got}{flag}")
    print()
    print("  Adding tellers adds zero correct increments.")
    print()


# ---------------------------------------------------------------- 3


def panel_lines():
    print("3. THE LOAD-BEARING LINE, AND THE INERT ONE")
    print(SEP)
    print("  The suspension moved by one position inside the same function.")
    print("  It is io_read either way, and it suspends either way:")
    for where, label in (("none", "no suspension at all      "),
                         ("before_read", "suspend(); v = bank[acct] "),
                         ("after_read", "v = bank[acct]; suspend() ")):
        got = bank(where, TELLERS, DEPOSITS)
        print(f"    {label} final={got:<5} lost={EXPECT - got:<4}"
              f" {'OK' if got == EXPECT else 'WRONG'}")
    print()
    print("  The inert line: colouring itself.")
    print(f"    stackful, nothing coloured       final="
          f"{bank('after_read', TELLERS, DEPOSITS)}")
    print(f"    stackless, all 4 frames coloured final="
          f"{run_green(rt.g3_teller)[0]}")
    print()
    print("  Identical. Four edits, no change in behaviour.")
    print()


# ---------------------------------------------------------------- 4


def panel_boundaries():
    print("4. BOUNDARIES -- where the effect vanishes")
    print(SEP)
    print("  (a) how many tellers it takes")
    for n in (1, 2, 3, 4, 8):
        exp, got = START + n * DEPOSITS, bank("after_read", n, DEPOSITS)
        print(f"    tellers={n:<2} final={got:<5} expected={exp:<5}"
              f" applied={got - START:<4} lost={exp - got}")
    print("    One teller has nothing to interleave with. The effect needs 2.")
    print()
    print("  (b) call depth D = frames above the suspending call, task body")
    print("      included. The bank workload is D=4.")
    print("           stackful          stackless, uncoloured        coloured")
    for d in range(5):
        _, sfe, sferr = run_chain(d, "stackful")
        _, _, uerr = run_chain(d, "uncoloured")
        _, ce, _ = run_chain(d, "coloured")
        sf = f"edits={sfe} {'ok' if sferr is None else sferr}"
        un = "runs, edits=0" if uerr is None else uerr.split(":")[0]
        print(f"      D={d}  {sf:<17} {un:<28} edits={ce}")
    print()
    print("    At D=0 the suspension is in the task's own body: zero frames")
    print("    are forced to change in either runtime, and the uncoloured")
    print("    stackless column runs too. The two runtimes are")
    print("    indistinguishable. Every tutorial example is written at D=0.")
    print()


# ---------------------------------------------------------------- 5


def coop_trial():
    """Real asyncio. The suspension point is a fixed location in the source."""
    acct = {"bal": START}

    async def io_read():
        v = acct["bal"]
        await asyncio.sleep(0)
        return v

    async def teller():
        for _ in range(DEPOSITS):
            acct["bal"] = await io_read() + 1

    async def main():
        await asyncio.gather(*[teller() for _ in range(TELLERS)])

    asyncio.run(main())
    return acct["bal"]


def preempt_trial(deposits, pad=0):
    """Real OS threads. The switch has to land in the window by luck."""
    acct = {"bal": START}

    def teller():
        for _ in range(deposits):
            b = acct["bal"]
            for _ in range(pad):
                pass
            acct["bal"] = b + 1

    ts = [threading.Thread(target=teller) for _ in range(TELLERS)]
    for t in ts:
        t.start()
    for t in ts:
        t.join()
    return acct["bal"]


def panel_certainty():
    print("5. CERTAIN vs RARE -- real asyncio against real threads")
    print(SEP)
    print(f"  A. the identical workload, 20 trials each, expected {EXPECT}")
    coop = [coop_trial() for _ in range(20)]
    pre = [preempt_trial(DEPOSITS) for _ in range(20)]
    print(f"    cooperative (asyncio)  values={sorted(set(coop))}"
          f"  wrong on {sum(x != EXPECT for x in coop)}/20 trials")
    print(f"    preemptive  (threads)  values={sorted(set(pre))}"
          f"  wrong on {sum(x != EXPECT for x in pre)}/20 trials")
    print()
    print("  B. how hard preemption has to be pushed to lose even one update")
    print("     (5 trials per row; `pad` widens the window with dead loops)")
    print("     deposits     pad    trials  wrong  worst loss")
    for deposits, pad in ((50, 0), (50_000, 0), (50_000, 20),
                          (200_000, 0), (200_000, 20)):
        vals = [preempt_trial(deposits, pad) for _ in range(5)]
        exp = START + TELLERS * deposits
        worst = max(exp - v for v in vals)
        print(f"     {deposits:<12,} {pad:<6} {5:<7} "
              f"{sum(v != exp for v in vals):<6} {worst:,}"
              f"  ({worst / (exp - START) * 100:.4f}% of increments)")
    print()
    print("  Cooperative scheduling did not make the race rarer. It made it")
    print("  the only behaviour the program has.")
    print()


# ---------------------------------------------------------------- 6


def panel_silence():
    print("6. THE LIMIT OF THE SAFETY NET")
    print(SEP)
    print("  'stackless refuses to run' holds only where the caller CONSUMES")
    print("  the value. Same un-coloured generator, three callers:")

    def io_read():
        v = {"bal": 7}
        yield
        return v

    log = []

    def uses():
        return io_read()["bal"] + 1

    def forwards():
        log.append(io_read())
        return "row logged"

    def truth_tests():
        return "found" if io_read() else "missing"

    for name, fn in (("uses it        (b['bal'] + 1)", uses),
                     ("forwards it    (log.append(b))", forwards),
                     ("truth-tests it (if b: ...)", truth_tests)):
        try:
            print(f"    {name:<31} -> {fn()!r}   NO ERROR")
        except Exception as e:
            print(f"    {name:<31} -> {type(e).__name__}: {e}")
    print(f"    what got logged in the forwarding case: "
          f"[{type(log[0]).__name__}], and it was never run")
    print()
    print("  Same three, against a real `async def` coroutine:")

    async def real_io_read():
        await asyncio.sleep(0)
        return {"bal": 7}

    with warnings.catch_warnings(record=True) as caught:
        warnings.simplefilter("always")
        try:
            c = real_io_read()
            print(f"    uses it                         -> {c['bal'] + 1!r}")
        except Exception as e:
            print(f"    uses it                         -> "
                  f"{type(e).__name__}: {e}")
        c2 = real_io_read()
        print(f"    truth-tests it (if user_exists) -> "
              f"{'found' if c2 else 'missing'!r}   NO ERROR")
        c3 = real_io_read()
        rlog = [c3]
        print(f"    forwards it                     -> "
              f"logged {type(rlog[0]).__name__}, NO ERROR")
        del c, c2, c3, rlog
        gc.collect()
    print("    the only thing Python said, and it said it at GC time:")
    for m in sorted({f"{w.category.__name__}: {w.message}" for w in caught}):
        print(f"      {m}")
    print()
    print("  bool() of an unstarted generator or an un-awaited coroutine is")
    print("  unconditionally True, so `if user_exists():` is always taken.")
    print()


# ---------------------------------------------------------------- 7

N_GREEN, N_THREADS = 10_000, 2_000


def rss_bytes():
    r = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
    return r if sys.platform == "darwin" else r * 1024


def measure_green(make):
    """Bytes of Python heap held by N_GREEN started-and-suspended tasks."""
    gc.collect()
    tracemalloc.start()
    base = tracemalloc.get_traced_memory()[0]
    tasks = [make() for _ in range(N_GREEN)]
    for t in tasks:
        t.send(None)                       # started and suspended: real state
    held = tracemalloc.get_traced_memory()[0] - base
    tracemalloc.stop()
    del tasks
    gc.collect()
    return held


def measure_threads():
    """RSS delta and spawn time for N_THREADS parked OS threads. A thread's
    stack is not a Python allocation, so tracemalloc cannot see it."""
    go = threading.Event()
    up = threading.Semaphore(0)

    def park():
        up.release()
        go.wait()

    gc.collect()
    before = rss_bytes()
    t0 = time.perf_counter()
    ths = [threading.Thread(target=park, daemon=True) for _ in range(N_THREADS)]
    for t in ths:
        t.start()
    for _ in range(N_THREADS):
        up.acquire()
    spawn = time.perf_counter() - t0
    held = rss_bytes() - before
    go.set()
    for t in ths:
        t.join()
    return held, spawn


def panel_cost():
    print("7. WHY ANYONE ACCEPTS COLOURING AT ALL")
    print(SEP)
    print("  ONE MACHINE, ONE RUN. These are the only wall-clock and memory")
    print("  numbers on the page; the banner above is their whole scope.")
    print()

    # Threads first: ru_maxrss is a high-water mark and never goes down, so
    # anything allocated before it would be baked into the baseline.
    tbytes, spawn = measure_threads()

    def one_frame():
        while True:
            yield
    flat = measure_green(one_frame)
    shared = {"acct": START}
    deep = measure_green(lambda: rt.g3_teller(shared, "acct", DEPOSITS))

    print("  state of suspended tasks")
    print(f"    green, {N_GREEN:,} one-frame tasks (tracemalloc)")
    print(f"      {flat:>14,} B   {flat / N_GREEN:>9.1f} B/task")
    print(f"    green, {N_GREEN:,} suspended g3_teller chains -- 5 generator")
    print(f"    frames each, because that is where the toy's tasks suspend")
    print(f"      {deep:>14,} B   {deep / N_GREEN:>9.1f} B/task"
          f"   = {deep / N_GREEN / 5:.1f} B/frame")
    print(f"    stackful, {N_THREADS:,} parked OS threads (RSS delta)")
    print(f"      {tbytes:>14,} B   {tbytes / N_THREADS:>9.1f} B/task")
    print(f"    ratio against the one-frame task: "
          f"{tbytes / N_THREADS / (flat / N_GREEN):.1f}x"
          f"    against the 5-frame chain: "
          f"{tbytes / N_THREADS / (deep / N_GREEN):.1f}x")
    print(f"    spawning the {N_THREADS:,} threads took {spawn * 1e3:.1f} ms"
          f" ({spawn / N_THREADS * 1e6:.1f} us each)")
    print()

    live = [one_frame() for _ in range(100_000)]
    for t in live:
        t.send(None)
    print("  ceiling")
    print(f"    green threads alive at once: {len(live):,} (no error)")
    del live
    gc.collect()

    stop = threading.Event()
    held, err = [], "no error"
    try:
        while len(held) < 60_000:
            t = threading.Thread(target=stop.wait, daemon=True)
            t.start()
            held.append(t)
    except BaseException as e:
        err = f"{type(e).__name__}: {e}"
    print(f"    OS threads alive at once:    {len(held):,}  -> {err}")
    stop.set()
    for t in held:
        t.join()
    print()
    print("  That is the price of OS threads, not the price of stackfulness --")
    print("  see the commentary's section 8.")


if __name__ == "__main__":
    banner()
    panel_headline()
    panel_arithmetic()
    panel_lines()
    panel_boundaries()
    panel_certainty()
    panel_silence()
    panel_cost()