"""Token budgets, sliding-window truncation, and what re-sending history costs.

Python standard library only. Run `python3 show_numbers.py` to print every
number the lesson quotes.
"""


def count_tokens(text: str) -> int:
    """Toy counter: one word (a run of non-space characters) counts as one token.

    Real tokenizers cut text into pieces they learned from data, so their counts
    differ from word counts, and differ from each other. To budget a real model,
    pass that model's own tokenizer as the `counter` argument of `truncate`.
    """
    return len(text.split())


def total_tokens(messages, counter=count_tokens, per_message=0) -> int:
    return sum(counter(m["content"]) + per_message for m in messages)


def truncate(messages, limit, keep_system=True, counter=count_tokens, per_message=0):
    """Sliding window: keep the newest messages that fit in `limit` tokens.

    messages     list of {"role": ..., "content": ...}, oldest first.
    keep_system  pin the leading system messages: they are always kept.
    per_message  extra tokens each message costs in the chat format (role
                 markers and the like). Take it from your model's template.

    The kept part is a run of the newest messages with no holes: the walk
    goes from newest to oldest and stops at the first message that does not
    fit. If the walk cut anything, the leading assistant replies of the run
    are then dropped, because a reply whose question was cut makes no sense
    on its own. Raises ValueError if the pinned messages alone exceed `limit`,
    or if no newest message is left, because answering without it is silent
    failure.
    """
    def cost(m):
        return counter(m["content"]) + per_message

    rest = list(messages)
    pinned = []
    if keep_system:
        n = 0
        while n < len(rest) and rest[n]["role"] == "system":
            n += 1
        pinned, rest = rest[:n], rest[n:]

    budget = limit - sum(cost(m) for m in pinned)
    if budget < 0:
        raise ValueError("the pinned system messages alone exceed the limit")

    kept = []
    for m in reversed(rest):
        if cost(m) > budget:
            break
        kept.append(m)
        budget -= cost(m)
    kept.reverse()
    if len(kept) < len(rest):                      # the walk cut something
        while kept and kept[0]["role"] == "assistant":
            kept.pop(0)                            # a reply whose question was cut
    if rest and not kept:
        raise ValueError("the newest message does not fit in the limit")
    return pinned + kept


def tokens_read(n_replies: int, tokens_per_message: int, window=None) -> int:
    """Total tokens the model reads while answering `n_replies` user messages.

    The app re-sends the history before every reply. Every message, yours or
    the bot's, has `tokens_per_message` tokens. With `window` set, the history
    is cut with `truncate` first. A reply's own tokens are not counted as read.
    """
    history, total = [], 0
    for _ in range(n_replies):
        history.append({"role": "user", "content": " ".join(["w"] * tokens_per_message)})
        sent = history if window is None else truncate(history, window)
        total += total_tokens(sent)
        history.append({"role": "assistant", "content": " ".join(["w"] * tokens_per_message)})
    return total


def attention_pairs(n: int) -> int:
    """Token pairs one full-attention head compares while reading an n-token
    prompt: token i is compared with the i tokens up to and including itself."""
    return n * (n + 1) // 2