"""A tiny model server you can read end to end: lesson 31 (build-a-tiny-model-server).

Three parts around a toy model: a cache of fixed-size pages (lesson 22), a batcher that steps every admitted chat together (lesson 24)
and metrics that report the waits (lesson 27). The "model" is a toy: a chat asks for a fixed number of tokens and each step gives every
admitted chat its next one, at the step time of lesson 24 (A100 80GB, Llama 3.1 8B). All of it runs in this one process on a simulated
clock, with no network code: the "load generator" is a list of send times, so nothing here can be pointed at a real service.
The widget (../widget.html) repeats `serve` and `report` in JavaScript; test_widget.py compares them. Standard library only.
"""
from __future__ import annotations

import importlib.util
import sys
from collections import deque, namedtuple
from pathlib import Path

# --- lesson 27's workload, percentiles and (through it) lesson 24's step time, loaded by path -----------------------------------

def _load_lesson_27():
    path = Path(__file__).resolve().parent.parent.parent / "why-average-speed-is-misleading" / "code" / "tail_sim.py"
    spec = importlib.util.spec_from_file_location("tail_sim_lesson27", path)
    module = importlib.util.module_from_spec(spec)
    sys.modules[spec.name] = module
    keep, sys.dont_write_bytecode = sys.dont_write_bytecode, True
    try:
        spec.loader.exec_module(module)
    finally:
        sys.dont_write_bytecode = keep
    return module


tail = _load_lesson_27()
step_seconds, percentile, intended_times = tail.step_seconds, tail.percentile, tail.intended_times

PAGE_TOKENS = 16                        # lesson 22 and lesson 28: tokens per cache page, 2 MiB of cache at 128 KiB a token
N_CHATS = 400                           # the first chats of lesson 27's seeded workload
WORKLOAD = {key: values[:N_CHATS] for key, values in tail.WORKLOAD.items()}
PAGES, LIMIT = 240, 16                  # the toy's cache (480 MiB; a real card holds 25,720 pages) and the most chats stepped together
LIGHT, HEAVY = 10, 40                   # chats sent per second


def pages_for(tokens: int) -> int:
    return -(-tokens // PAGE_TOKENS)


# --- the cache ----------------------------------------------------------------------------------------------------------------

class Pages:
    """Lesson 22's free list and page tables, as counts only: no keys and values are stored."""

    def __init__(self, total: int):
        self.total, self.free, self.tables = total, deque(range(total)), {}

    def grow(self, chat: int, tokens: int) -> None:
        """A chat holding `tokens` tokens needs enough pages for them: it takes one more whenever its last page is full."""
        table = self.tables.setdefault(chat, [])
        while len(table) * PAGE_TOKENS < tokens:
            table.append(self.free.popleft())

    def release(self, chat: int) -> None:
        self.free.extend(self.tables.pop(chat))

    def used(self) -> int:
        return self.total - len(self.free)


# --- the batcher ------------------------------------------------------------------------------------------------------------

# One row per step: when it started, how long it took, chats that got a token, of those how many just joined, chats still waiting,
# pages in use, pages promised but not yet taken, and what kept the next chat out ("cache", "batch", "cache and batch" or "" for nothing).
Step = namedtuple("Step", "start seconds stepped joined queue held owed blocked")
Run = namedtuple("Run", "first done steps")      # when each chat's first and last token came out, and every step


def serve(workload: dict, rate: float, pages: int = PAGES, limit: int = LIMIT) -> Run:
    prompts, replies, arrive = workload["prompts"], workload["replies"], intended_times(workload["gaps"], rate)
    n = len(prompts)
    assert pages >= pages_for(max(p + r for p, r in zip(prompts, replies))), "one chat needs more pages than the cache has"
    cache, queue, running, made = Pages(pages), deque(), [], [0] * n
    first, done, steps = [0.0] * n, [0.0] * n, []
    promised, now, sent, finished = 0, 0.0, 0, 0
    while finished < n:
        if not queue and not running and arrive[sent] > now:
            now = arrive[sent]                                          # nothing to do: wait for the next chat
        while sent < n and arrive[sent] <= now:
            queue.append(sent)                                          # 1. a request arrives and waits in the queue
            sent += 1
        joining, blocked = [], ""
        while queue:                                                    # 2. admit it when a place and all its pages are free
            need = pages_for(prompts[queue[0]] + replies[queue[0]])
            no_place, no_pages = len(running) + len(joining) == limit, promised + need > pages
            blocked = " and ".join(wall for wall, hit in (("cache", no_pages), ("batch", no_place)) if hit)
            if blocked:
                break
            promised += need
            joining.append(queue.popleft())
        seconds = step_seconds([prompts[i] + made[i] for i in running], sum(prompts[i] for i in joining))      # 3. one step, all together
        start, now = now, now + seconds
        running += joining
        for i in running:
            made[i] += 1
            cache.grow(i, prompts[i] + made[i])                         # 4. a page more when the last one is full
        for i in joining:
            first[i] = now
        stepped, ended = len(running), [i for i in running if made[i] == replies[i]]
        for i in ended:
            done[i] = now
            cache.release(i)                                            # 4. all its pages back when the chat ends
            promised -= pages_for(prompts[i] + replies[i])
            running.remove(i)
        finished += len(ended)
        steps.append(Step(start, seconds, stepped, len(joining), len(queue), cache.used(), promised - cache.used(), blocked))
    return Run(first, done, steps)


# --- the metrics ------------------------------------------------------------------------------------------------------------

def report(workload: dict, rate: float, run: Run) -> dict:
    """5. What the manager's clipboard shows: time to first token and time between tokens (every gap of every chat, pooled; p50, p99), the longest queue, the most pages used."""
    arrive = intended_times(workload["gaps"], rate)
    ttft = [f - a for f, a in zip(run.first, arrive)]
    tbt = [s.seconds for s in run.steps for _ in range(s.stepped - s.joined)]       # a running chat gets a word every step: its gap is the step's length
    cache_full = [s.start for s in run.steps if "cache" in s.blocked]
    return {"ttft_p50": percentile(ttft, 50), "ttft_p99": percentile(ttft, 99), "tbt_p50": percentile(tbt, 50), "tbt_p99": percentile(tbt, 99),
            "queue_max": max(s.queue for s in run.steps), "pages_max": max(s.held for s in run.steps),
            "cache_full_at": cache_full[0] if cache_full else None, "served": len(ttft) / max(run.done),
            "cache_steps": len(cache_full), "batch_steps": sum(1 for s in run.steps if "batch" in s.blocked)}


def load_test(rate: float, pages: int = PAGES, limit: int = LIMIT, workload: dict = WORKLOAD) -> dict:
    """6. The in-process load generator: chats arrive at random times that average `rate` a second, and the server reports."""
    return report(workload, rate, serve(workload, rate, pages, limit))


MIN_QUEUE = 4                           # a wall is named only if this many chats (1 in 100 of the 400) waited at once


def bottleneck(metrics: dict) -> str:
    """Which wall held the next chat out in more steps: "cache" (memory binds), "batch" (the batch limit binds), "both" (neither more than twice
    as often as the other) or "none" (fewer than MIN_QUEUE chats ever waited at once)."""
    c, b = metrics["cache_steps"], metrics["batch_steps"]
    return "none" if metrics["queue_max"] < MIN_QUEUE else "both" if 2 * min(c, b) >= max(c, b) else "cache" if c > b else "batch"