"""Prompt caching as a small simulation: which start of a prompt is reused, and what that does to the bill.

Python 3 standard library only (plus three earlier lessons' code, imported, not copied). Run `python3 show_numbers.py`
to print every number the lesson quotes.

The model, in four rules:
  1. A request's prompt is a list of tokens (lesson 1's tokenizer, `bpe.encode`).
  2. The provider keeps a start of a request only up to a mark the caller set, or, with no mark, to the end of the request.
     A new request reuses a kept start only if it begins with the whole of it; it reads all of that start and nothing
     shorter, so one changed token anywhere inside the kept start loses all of it.
  3. A rule (minimum length, lifetime of an unused copy, three prices) decides what each token costs: tokens read from
     a kept start, tokens written to a new one, tokens neither.
  4. A kept start is refreshed each time a request reads it and is gone once it has gone unused for the lifetime.

Prices are dollars per million tokens. Money is counted in whole units of 1/1000 of a micro-dollar per token
(a price of $2 per million is 2,000 units per token), so every bill is an exact integer; the widget does the same sums.
The rules and prices below are read from provider pages dated 2026-10-19 (see the lesson's Further reading). They change.
"""
from __future__ import annotations

import functools
import math
import sys
from dataclasses import dataclass
from decimal import Decimal
from fractions import Fraction
from pathlib import Path

LESSONS = Path(__file__).resolve().parents[2]
for _lesson in ("what-is-a-token", "why-ai-bill-is-high", "why-ai-waits-before-the-first-word"):
    sys.path.append(str(LESSONS / _lesson / "code"))     # appended, so this folder's own modules win a name clash
import bill as chat_bill                                  # lesson 7: the chat and its bill without caching
import bpe                                                # lesson 1: the tokenizer
import prefill_sim                                        # lesson 23: the best-case time of reading a prompt

UNITS = 1000                    # units per (dollar per million tokens)


# --- the rules -----------------------------------------------------------------------------------

@dataclass(frozen=True)
class Rule:
    key: str                    # the letter the widget and the lesson use
    label: str                  # who published it
    price: Decimal              # normal input price, dollars per million tokens
    write: Decimal              # price of a token written to the copy
    read: Decimal               # price of a token read from the copy
    output: Decimal             # output price, dollars per million tokens
    minimum: int                # the cached start must be at least this many tokens
    lifetime: int               # minutes an unused copy lives; every read starts the count again

    def units(self, kind: str) -> int:
        return int(getattr(self, kind) * UNITS)


D = Decimal
# Anthropic: platform.claude.com/docs/en/about-claude/pricing (prices) and .../build-with-claude/prompt-caching (minimum, lifetime). Read 2026-10-19.
# OpenAI: developers.openai.com/api/docs/pricing (prices) and .../guides/prompt-caching (minimum, lifetime). Read 2026-10-19.
RULES = [
    Rule("A", "Anthropic, Sonnet 5.5, 5-minute cache", D("2"), D("2.50"), D("0.10"), D("10"), minimum=512, lifetime=5),
    Rule("B", "Anthropic, Sonnet 5.5, 1-hour cache", D("2"), D("4"), D("0.10"), D("10"), minimum=512, lifetime=60),
    Rule("C", "Anthropic, Haiku 5.5, 5-minute cache", D("0.10"), D("0.125"), D("0.01"), D("0.50"), minimum=512, lifetime=5),
    Rule("D", "OpenAI, gpt-6.1-sol", D("2.00"), D("2.50"), D("0.10"), D("10"), minimum=1024, lifetime=30),
]
RULE_BY_KEY = {r.key: r for r in RULES}


def break_even_reads(rule: Rule) -> int:
    """Fewest reads that pay back one write: w + r*h <= 1 + r, so r >= (w - 1) / (1 - h), rounded up. Exact, from the units."""
    w, h, p = rule.units("write"), rule.units("read"), rule.units("price")
    return max(0, math.ceil(Fraction(w - p, p - h)))


# --- the match -----------------------------------------------------------------------------------

def shared_start(old, new) -> int:
    """How many tokens, counting from the first, the two prompts have in common before the first difference."""
    n = 0
    for a, b in zip(old, new):
        if a != b:
            break
        n += 1
    return n


def reuse(stored, new) -> int:
    """Tokens a new request reads from one kept start: all of it when the request begins with all of it, else none."""
    return len(stored) if shared_start(stored, new) == len(stored) else 0


@dataclass(frozen=True)
class Billed:
    read: int                   # tokens read from a stored copy
    written: int                # tokens written to a new copy
    plain: int                  # tokens billed at the normal price
    units: int                  # the bill for the input, in units
    plain_units: int            # the same prompt with no caching, in units


def bill_request(rule: Rule, total: int, reused: int, mark: int | None = None) -> Billed:
    """Bill one prompt of `total` tokens that begins with a live kept start of `reused` tokens (0 when it begins with none).

    Reading covers the whole kept start, and only if it reaches the minimum. Writing covers the rest of the part the caller
    marked (`mark` tokens; the whole prompt when None) and happens only if that part reaches the minimum.
    """
    read = reused if reused >= rule.minimum else 0
    stored = total if mark is None else min(mark, total)
    written = stored - read if stored >= rule.minimum and stored > read else 0
    plain = total - read - written
    units = read * rule.units("read") + written * rule.units("write") + plain * rule.units("price")
    return Billed(read, written, plain, units, total * rule.units("price"))


class Cache:
    """The starts a provider keeps for one rule. `request` bills one prompt arriving at `minute` and keeps what it stores."""

    def __init__(self, rule: Rule):
        self.rule = rule
        self.copies: list[list] = []         # [tokens, minute last used]

    def request(self, prompt, minute: float, mark: int | None = None) -> Billed:
        live = [c for c in self.copies if minute - c[1] <= self.rule.lifetime]
        best = max(live, key=lambda c: reuse(c[0], prompt), default=None)
        billed = bill_request(self.rule, len(prompt), reuse(best[0], prompt) if best else 0, mark)
        if billed.read:
            best[1] = minute                  # a read refreshes the copy it used
        if billed.written:
            live.append([list(prompt[:billed.read + billed.written]), minute])
        self.copies = live
        return billed


# --- a chat that re-sends its history ---------------------------------------------------------------

def chat_prompts(chat, stamp: bool = False):
    """The prompt of each turn of a lesson 7 chat: system prompt, every earlier message and reply, then the new message.
    Tokens are labels, equal where the text is equal. With stamp=True a clock line sits first and changes every turn."""
    if chat.reasoning or chat.keep is not None:
        raise ValueError("the chat must send its whole history and no hidden reasoning")
    system = [("system", i) for i in range(chat.system)]
    history: list = []
    for k in range(chat.turns):
        message = [("user", k, i) for i in range(chat.message)]
        yield ([("clock", k)] if stamp else []) + system + history + message
        history += message + [("reply", k, i) for i in range(chat.reply)]


def chat_cost(chat, rule: Rule, gap: float = 1, stamp: bool = False, cached: bool = True) -> dict:
    """The units a whole chat costs: input, with or without caching, plus output. `gap` is the minutes between turns."""
    cache, read, written, plain, units = Cache(rule), 0, 0, 0, 0
    for k, prompt in enumerate(chat_prompts(chat, stamp)):
        b = cache.request(prompt, k * gap) if cached else bill_request(rule, len(prompt), 0, 0)
        read, written, plain, units = read + b.read, written + b.written, plain + b.plain, units + b.units
    out = chat.turns * chat.reply * rule.units("output")
    return {"read": read, "written": written, "plain": plain, "input": units, "output": out, "total": units + out}


def dollars(units: int) -> Decimal:
    return Decimal(units) / Decimal(10**9)


def percent_less(base: int, cost: int) -> int:
    """Whole percent by which cost is below base (negative when above), rounded half up, in integers."""
    diff = base - cost
    return (200 * diff + base) // (2 * base) if diff >= 0 else -((200 * -diff + base) // (2 * base))


# --- time: lesson 23's best case ------------------------------------------------------------------------

def first_word_seconds(total: int, read: int) -> tuple[float, float]:
    """Best-case reading time at the chip's peak, without and with a reused start of `read` tokens (lesson 23's model).
    The newest token is always computed, so at least one token is read."""
    return prefill_sim.prefill(total).seconds, prefill_sim.prefill(max(total - read, 1)).seconds


# --- the widget's view -------------------------------------------------------------------------------------

MARKER = "[document]"
DOC_BASE = 1_000_000            # ids of the stand-in document tokens; real ids stay below 768
DOC_LENGTHS = [0, 400, 1200, 5000]
DEFAULT_DOC = 1200
AGES = [1, 4, 6, 40, 120]       # minutes since the earlier request
DEFAULT_AGE = 1
FIXED = "You are the help bot of Harbour Bakery. Answer in two short sentences.\n" + MARKER + "\nQuestion: "
EARLIER = FIXED + "Do you open on Sundays?"
CLOCK = "Time now: Tue 20 Oct, 14:05"
PRESETS = [
    {"name": "New question", "text": FIXED + "Is the rye bread vegan?"},
    {"name": "Date at the top", "text": CLOCK + "\n" + EARLIER},
    {"name": "Date at the end", "text": EARLIER + "\n" + CLOCK},
]
OPENING = PRESETS[0]["text"]


@functools.cache
def merges() -> list[tuple[int, int]]:
    """Lesson 1's 512 learned merges (trained on its own 15 KB corpus, about 0.8 s)."""
    return bpe.train(bpe.load_corpus())


def tokens(text: str, doc: int, learned=None) -> list[int]:
    """Token ids of a prompt. Each [document] marker stands for `doc` stand-in tokens, the same ones wherever the text matches."""
    learned = merges() if learned is None else learned
    out: list[int] = []
    for j, part in enumerate(text.split(MARKER)):
        if j:
            out.extend(DOC_BASE + (j - 1) * doc + i for i in range(doc))
        out.extend(bpe.encode(part, learned))
    return out


def _show(raw: bytes) -> str:
    try:
        text = raw.decode("utf-8")
    except UnicodeDecodeError:
        return " ".join(f"{b:02X}" for b in raw)
    return text.replace("\n", "↵")


def segments(ids: list[int], matched: int, learned=None) -> list[list]:
    """What the widget draws: [label, 'match' | 'rest', 1 for a document block] per token, with a run of stand-in document tokens drawn as one block."""
    table = bpe.vocab_bytes(merges() if learned is None else learned)
    out: list[list] = []                      # [label, mark, document tokens in the run (0 for an ordinary token)]
    for i, t in enumerate(ids):
        mark = "match" if i < matched else "rest"
        if t < DOC_BASE:
            out.append([_show(table[t]), mark, 0])
        elif out and out[-1][2] and out[-1][1] == mark:
            out[-1][2] += 1
        else:
            out.append(["", mark, 1])
    return [[f"document, {commas(n)} tokens", mark, 1] if n else [label, mark, 0] for label, mark, n in out]


def commas(n: int) -> str:
    return f"{n:,}"


def money(units: int) -> str:
    """Dollars to the millionth, rounded half up in integers so that the widget's JavaScript prints the same digits."""
    micro = (units + 500) // 1000
    return f"${micro // 10**6}.{micro % 10**6:06d}"


def summary(rule: Rule) -> str:
    """The rule in one line for the widget's list."""
    return f"{rule.key}: write {float(rule.write / rule.price):g}x, read {float(rule.read / rule.price):g}x, {rule.lifetime} min"


def rate(x: Decimal) -> str:
    s = f"${float(x):.3f}"
    return s[:-1] if s.endswith("0") else s


def mark_at(text: str) -> int:
    """Where the widget puts the mark: just after the line holding [document], else just before the last line."""
    i = text.find(MARKER)
    if i >= 0:
        j = text.find("\n", i)
        return len(text) if j < 0 else j + 1
    return text.rfind("\n") + 1


def mark_tokens(text: str, doc: int, learned=None) -> int:
    """How many tokens of the prompt come before the mark."""
    return len(tokens(text[:mark_at(text)], doc, learned))


def view(old: str, new: str, rule: Rule, doc: int, age: float, marked: bool, learned=None) -> dict:
    """Every line of words the widget shows for one earlier request, one new request and the settings.
    With `marked`, each request stores its start up to the mark; without it, only the whole request."""
    a, b = tokens(old, doc, learned), tokens(new, doc, learned)
    total = len(b)
    kept = mark_tokens(old, doc, learned) if marked else len(a)
    stored = kept if kept >= rule.minimum else 0
    alive = age <= rule.lifetime
    bl = bill_request(rule, total, reuse(a[:stored], b) if alive else 0, mark_tokens(new, doc, learned) if marked else None)
    d = shared_start(a[:stored], b)
    out = {"old": segments(a, stored, learned), "new": segments(b, bl.read, learned), "oldCount": len(a), "newCount": total}
    if not total:
        out.update(match="Type a new request to compare.", billed="", cost="", why="")
        return out
    if bl.read:
        out["match"] = f"{commas(bl.read)} of {commas(total)} tokens are reused from the stored start."
    elif stored and alive and d and not marked:
        out["match"] = "No token is reused: without a mark, only the whole earlier request is stored."
    else:
        out["match"] = "No token is reused."
    out["billed"] = (f"Billed: {commas(bl.read)} tokens at {rate(rule.read)} (read), {commas(bl.written)} at {rate(rule.write)} (write), "
                     f"{commas(bl.plain)} at {rate(rule.price)} (normal), per million.")
    pct = percent_less(bl.plain_units, bl.units)
    change = "about the same" if pct == 0 else f"{pct}% less" if pct > 0 else f"{-pct}% more"
    out["cost"] = f"Input cost: {money(bl.units)} with caching, {money(bl.plain_units)} without: {change}."
    if total < rule.minimum:
        out["why"] = f"The prompt is shorter than the {commas(rule.minimum)}-token minimum, so nothing is stored or read."
    elif not stored:
        out["why"] = (f"The earlier request stores {commas(kept)} tokens, under the {commas(rule.minimum)}-token minimum, "
                      "so there is no stored start to reuse.")
    elif not alive:
        out["why"] = f"The earlier request is older than {rule.lifetime} minutes, so its stored start is gone and nothing is read."
    elif bl.read:
        left = f"The other {commas(total - bl.read)} tokens are read again." if total > bl.read else "Nothing is left to read again."
        out["why"] = f"The new request begins with the whole stored start, so it is read from the cache. {left}"
    elif not d:
        out["why"] = "The very first token differs from the stored start, so none of it is reused."
    else:
        ends = "is shorter than the stored start" if d == total else f"stops matching the stored start after {commas(d)} of its {commas(total)} tokens"
        if marked:
            out["why"] = f"The new request {ends}. A stored start is reused whole or not at all."
        else:
            out["why"] = ("Nothing is reused, and the whole prompt is written as a new copy. "
                          f"Without a mark, only the whole earlier request is stored, and the new request {ends}.")
    return out