#!/usr/bin/env python3
"""Prompt-caching cost demo: same questions over one large document, with and
without prompt caching, priced from the response usage fields.

What it shows
-------------
A repeated-context workload (one big document, many questions) is exactly the
shape prompt caching is built for. Run 1 sends the document as ordinary input
on every call (full price each time). Run 2 marks the document with a cache
breakpoint: the first call pays a one-time *write* premium, every later call
pays the cheap *read* rate. The script prints token usage and dollar cost for
both so the savings are measured, not asserted.

How it works (the mechanic the guide teaches)
---------------------------------------------
The document is the *stable prefix* (in `system`); the question is the *volatile
tail* (in `messages`). Caching is a prefix match, so the document must be
byte-identical across calls and the question must come after the cached block.
Calls run sequentially: a cache entry is only readable once the first response
has come back, so parallel calls would all miss.

Run it
------
  export ANTHROPIC_API_KEY=sk-...      # or use `ant auth login`
  pip install anthropic
  python caching_cost_demo.py                       # default Gutenberg text
  python caching_cost_demo.py path/to/local.txt     # your own document
  python caching_cost_demo.py https://example/x.txt # any plain-text URL

Cost: with the defaults below (~15k-token doc, 3 questions, Opus 4.8) a full
run is well under $1. Set MODEL = "claude-haiku-4-5" to run it ~5x cheaper.
Paste the printed output back and it becomes dated evidence for the guide.
"""

from __future__ import annotations

import sys
import urllib.request

import anthropic

# --- Configuration ----------------------------------------------------------

# House default is the most capable model. Swap to "claude-haiku-4-5" to run
# this demo ~5x cheaper; the caching mechanics and the shape of the result are
# identical, only the per-token price differs.
MODEL = "claude-opus-4-8"

# Per-MTok list prices (USD). Verify against the current Anthropic pricing page
# before using any number from this run in the published guide, and DATE it.
#   input:  base price for uncached input tokens
#   output: base price for output tokens
#   cache_write_multiplier: 1.25x input for the 5-minute TTL (2.0x for 1h)
#   cache_read_multiplier:  0.10x input
PRICING = {
    "claude-opus-4-8": {"input": 5.00, "output": 25.00},
    "claude-sonnet-4-6": {"input": 3.00, "output": 15.00},
    "claude-haiku-4-5": {"input": 1.00, "output": 5.00},
}
CACHE_WRITE_MULTIPLIER = 1.25  # 5-minute TTL
CACHE_READ_MULTIPLIER = 0.10

# Keep the document comfortably above the model's minimum cacheable prefix
# (4096 tokens on Opus 4.8) while bounding cost. ~60k chars is ~15k tokens.
MAX_DOC_CHARS = 60_000

# A public-domain text (Project Gutenberg). Override with a file path or URL as
# the first CLI argument.
DEFAULT_DOC_URL = "https://www.gutenberg.org/cache/epub/1342/pg1342.txt"

QUESTIONS = [
    "In one sentence, what is the central conflict of this text?",
    "Name two of the main characters and describe each in one line.",
    "Quote one sentence from the text that reflects its central theme.",
]

MAX_TOKENS = 300  # answers are short; this keeps output cost small and focused


# --- Document loading -------------------------------------------------------


def load_document(source: str) -> str:
    """Load the document from a URL or local path, and trim boilerplate/size."""
    if source.startswith(("http://", "https://")):
        req = urllib.request.Request(source, headers={"User-Agent": "cost-demo/1.0"})
        with urllib.request.urlopen(req, timeout=30) as resp:
            text = resp.read().decode("utf-8", errors="replace")
    else:
        with open(source, "r", encoding="utf-8", errors="replace") as fh:
            text = fh.read()

    # Strip Project Gutenberg's header/footer boilerplate if present, so the
    # cached prefix is the actual work, not licensing text.
    start = text.find("*** START OF")
    if start != -1:
        text = text[text.find("\n", start) + 1 :]
    end = text.find("*** END OF")
    if end != -1:
        text = text[:end]

    text = text.strip()
    if len(text) > MAX_DOC_CHARS:
        text = text[:MAX_DOC_CHARS]
    return text


# --- Cost accounting --------------------------------------------------------


def price_usage(usage, model: str) -> dict[str, float]:
    """Turn a response `usage` object into a per-call cost breakdown (USD)."""
    rates = PRICING[model]
    in_rate = rates["input"] / 1_000_000
    out_rate = rates["output"] / 1_000_000

    # These fields are the source of truth. cache_read_input_tokens > 0 proves a
    # hit; if it stays 0 across the cached run, a silent invalidator is at work.
    uncached_in = usage.input_tokens
    cache_write = getattr(usage, "cache_creation_input_tokens", 0) or 0
    cache_read = getattr(usage, "cache_read_input_tokens", 0) or 0
    output = usage.output_tokens

    return {
        "uncached_input_tokens": uncached_in,
        "cache_write_tokens": cache_write,
        "cache_read_tokens": cache_read,
        "output_tokens": output,
        "cost": (
            uncached_in * in_rate
            + cache_write * in_rate * CACHE_WRITE_MULTIPLIER
            + cache_read * in_rate * CACHE_READ_MULTIPLIER
            + output * out_rate
        ),
    }


# --- The two runs -----------------------------------------------------------


def run(client, document: str, use_cache: bool) -> list[dict[str, float]]:
    """Ask every question over the document, sequentially. With use_cache, the
    document carries a cache breakpoint so the first call writes and the rest
    read; without it, every call pays full price for the document."""
    # The document is the stable prefix. Only add cache_control in the caching
    # run; the question always goes in `messages`, after the cached block.
    system_block = {"type": "text", "text": document}
    if use_cache:
        system_block["cache_control"] = {"type": "ephemeral"}

    results = []
    for question in QUESTIONS:
        resp = client.messages.create(
            model=MODEL,
            max_tokens=MAX_TOKENS,
            system=[system_block],
            messages=[{"role": "user", "content": question}],
        )
        results.append(price_usage(resp.usage, MODEL))
    return results


# --- Reporting --------------------------------------------------------------


def report(label: str, rows: list[dict[str, float]]) -> float:
    total = sum(r["cost"] for r in rows)
    print(f"\n{label}")
    print(
        f"  {'#':>2}  {'uncached_in':>11}  {'cache_write':>11}  "
        f"{'cache_read':>10}  {'output':>6}  {'cost($)':>9}"
    )
    for i, r in enumerate(rows, 1):
        print(
            f"  {i:>2}  {r['uncached_input_tokens']:>11}  "
            f"{r['cache_write_tokens']:>11}  {r['cache_read_tokens']:>10}  "
            f"{r['output_tokens']:>6}  {r['cost']:>9.4f}"
        )
    print(f"  {'':>2}  {'':>11}  {'':>11}  {'':>10}  {'total':>6}  {total:>9.4f}")
    return total


def main() -> None:
    source = sys.argv[1] if len(sys.argv) > 1 else DEFAULT_DOC_URL
    if MODEL not in PRICING:
        sys.exit(f"No pricing configured for {MODEL}; add it to PRICING.")

    document = load_document(source)
    client = anthropic.Anthropic()

    approx_tokens = client.messages.count_tokens(
        model=MODEL, messages=[{"role": "user", "content": document}]
    ).input_tokens

    print(f"Model:            {MODEL}")
    print(f"Document source:  {source}")
    print(f"Document size:    {len(document):,} chars (~{approx_tokens:,} tokens)")
    print(f"Questions:        {len(QUESTIONS)}")
    if approx_tokens < 4096:
        print(
            "WARNING: document is below the 4096-token minimum cacheable "
            "prefix for Opus 4.8 — caching will not engage."
        )

    without = report(
        "WITHOUT caching (document re-sent at full price each call):",
        run(client, document, use_cache=False),
    )
    with_ = report(
        "WITH caching (document cached: 1 write, then cheap reads):",
        run(client, document, use_cache=True),
    )

    saved = without - with_
    pct = (saved / without * 100) if without else 0.0
    print("\nSummary")
    print(f"  without caching: ${without:.4f}")
    print(f"  with caching:    ${with_:.4f}")
    print(f"  saved:           ${saved:.4f}  ({pct:.1f}% lower)")
    print(
        "\nNote: savings scale with reuse. More questions over the same cached "
        "document widen the gap; a single question would not repay the write "
        "premium. That is the guide's thesis, measured."
    )


if __name__ == "__main__":
    main()