"""
What it costs to fill a model's context window once: input side only, one request.

For each of the 13 chat models in ../pricing.json this script computes the price of
a single request whose prompt is 100%, 50% and 10% of the model's maximum input
size, plus the price of a fixed 200,000-token prompt. Four models carry a
long-context surcharge (a higher rate once the prompt is over a threshold); those
rates are not in pricing.json, so they are written down in TIERS below.

Everything else comes from ../pricing.json (site/public/data/pricing.json) at the
time the script is run. "context" there is the LiteLLM dataset's max_input_tokens,
an input limit, not the output limit.

Method (all constants are at the top of the file):
  tokens = round(context * fraction)
  rate   = the model's higher rate if it has one and tokens > threshold, else the base rate
  cost   = tokens * rate / 1,000,000          (USD, rounded to 4 places when written)
Ratios and counts such as "38.1x" or "fills per $100" are computed from the unrounded
costs, so they can differ in the last digit from a ratio of two rounded table cells.

What this does NOT model: output tokens, reasoning tokens, prompt caching, batch,
priority or flex tiers, regional processing surcharges, or whether a model stays
accurate on a full-size prompt.

Run:
    python context_fill_cost.py            # checks, then rewrite context_fill_cost.results.json
    python context_fill_cost.py --html     # also print table bodies and the SVG used in the article

If a base rate in pricing.json no longer matches the one recorded in TIERS the
script refuses to write anything: the surcharge data has to be re-read first.
"""

import json
import sys
from decimal import Decimal
from pathlib import Path

# ---- assumptions (edit these to model something else) ----
FILL_FRACTIONS = {"full": 1.0, "half": 0.5, "tenth": 0.1}
COMMON_TOKENS = 200_000      # the largest prompt that every tracked model accepts
BUDGET_USD = 100.0           # used for "how many full fills does $100 buy"
LONG_CONTEXT_MIN = 900_000   # max input at or above this counts as "long context"
PRICEY_THRESHOLD_USD = 1.0   # "pricey" = one full fill costs at least this much

# Long-context surcharges. Not in pricing.json: the site's daily sync keeps only the
# flat rate. Read from the LiteLLM dataset (download MD5 7d541b8e2238009e32ce7df21f07c1be)
# on 2026-10-13; the field names are given so anyone can look them up.
# Whole-prompt rule as implemented in LiteLLM
# litellm/litellm_core_utils/llm_cost_calc/utils.py at commit
# c2bc5aecff15604f42180d7725c8b34c540bfa11: once the prompt token count is strictly
# greater than the threshold, the input AND output rates for that request switch to
# the higher ones, and the whole prompt is billed at the higher input rate.
# The vendor's own bill is the authority; neither vendor page says in words whether
# the higher rate covers the whole prompt or only the part above the threshold.
TIERS_READ_ON = "2026-10-13"
TIERS_DATASET_URL = "https://raw.githubusercontent.com/BerriAI/litellm/main/model_prices_and_context_window.json"
TIERS_DATASET_MD5 = "7d541b8e2238009e32ce7df21f07c1be"
TIERS_CALC_COMMIT = "c2bc5aecff15604f42180d7725c8b34c540bfa11"
TIERS = {
    "gpt-5.6": {
        "above": 272_000, "in_base": 4.0, "in_above": 8.0, "out_base": 20.0, "out_above": 30.0,
        "in_field": "input_cost_per_token_above_272k_tokens",
        "out_field": "output_cost_per_token_above_272k_tokens",
    },
    "gpt-5.6-terra": {
        "above": 272_000, "in_base": 2.0, "in_above": 4.0, "out_base": 12.0, "out_above": 18.0,
        "in_field": "input_cost_per_token_above_272k_tokens",
        "out_field": "output_cost_per_token_above_272k_tokens",
    },
    "gpt-5.6-luna": {
        "above": 272_000, "in_base": 0.2, "in_above": 0.4, "out_base": 1.2, "out_above": 1.8,
        "in_field": "input_cost_per_token_above_272k_tokens",
        "out_field": "output_cost_per_token_above_272k_tokens",
    },
    "gemini/gemini-3.1-pro-preview": {
        "above": 200_000, "in_base": 2.0, "in_above": 4.0, "out_base": 12.0, "out_above": 18.0,
        "in_field": "input_cost_per_token_above_200k_tokens",
        "out_field": "output_cost_per_token_above_200k_tokens",
    },
}

PRICING_PATH = Path(__file__).resolve().parent.parent / "pricing.json"
RESULTS_PATH = Path(__file__).with_name("context_fill_cost.results.json")
VENDOR_COLORS = {"Anthropic": "#2563eb", "OpenAI": "#15803d", "Google": "#6b7280"}


# ---------------------------------------------------------------- computation
def rate_for(model, tokens):
    """Input rate ($/1M) that applies to a prompt of `tokens` (strict > threshold)."""
    t = TIERS.get(model["id"])
    if t and tokens > t["above"]:
        return t["in_above"]
    return model["in"]


def build_models(pricing):
    """Returns (models for results.json, raw unrounded costs keyed by id)."""
    models, raw = [], {}
    for m in pricing["chat"]:
        t = TIERS.get(m["id"])
        fill, rawfill = {}, {}
        for name, frac in FILL_FRACTIONS.items():
            tokens = round(m["context"] * frac)
            flat = tokens * m["in"] / 1e6
            appl = tokens * rate_for(m, tokens) / 1e6
            rawfill[name] = {"flat": flat, "appl": appl}
            fill[name] = {
                "tokens": tokens,
                "cost_flat": round(flat, 4),
                "cost_applicable": round(appl, 4),
                "tier_applied": bool(t and tokens > t["above"]),
            }
        c200 = COMMON_TOKENS * rate_for(m, COMMON_TOKENS) / 1e6
        full = rawfill["full"]["appl"]
        long_ctx = m["context"] >= LONG_CONTEXT_MIN
        models.append({
            "id": m["id"], "name": m["name"], "vendor": m["vendor"],
            "context": m["context"], "in": m["in"], "out": m["out"],
            "tier": ({k: t[k] for k in ("above", "in_base", "in_above", "out_base", "out_above",
                                          "in_field", "out_field")} if t else None),
            "fill": fill,
            "cost_200k": round(c200, 4),
            "requests_per_100usd_full": round(BUDGET_USD / full, 1),
            "full_over_tenth": round(full / rawfill["tenth"]["appl"], 1),
            "long_context": long_ctx,
            "fits_but_pricey": bool(long_ctx and full >= PRICEY_THRESHOLD_USD),
        })
        raw[m["id"]] = {"fill": rawfill, "c200": c200}
    order = sorted(models, key=lambda x: (-raw[x["id"]]["fill"]["full"]["appl"], x["id"]))
    return order, raw


def build_summary(models, raw):
    by_id = {m["id"]: m for m in models}
    full = lambda i: raw[i]["fill"]["full"]["appl"]  # noqa: E731  (unrounded)
    ids = [m["id"] for m in models]                  # already sorted by full cost, high to low
    hi, lo = ids[0], ids[-1]
    long_ids = [i for i in ids if by_id[i]["long_context"]]
    lhi, llo = long_ids[0], long_ids[-1]
    asc = sorted(long_ids, key=lambda i: (full(i), i))
    steps = [{
        "low_id": a, "low_cost": round(full(a), 4),
        "high_id": b, "high_cost": round(full(b), 4),
        "ratio": round(full(b) / full(a), 4),
    } for a, b in zip(asc, asc[1:])]
    steps.sort(key=lambda s: -s["ratio"])
    c200_min = min(raw[i]["c200"] for i in ids)
    c200_max = max(raw[i]["c200"] for i in ids)
    tiered = [i for i in ids if by_id[i]["tier"]]
    free = [f for f in FILL_FRACTIONS if not any(by_id[i]["fill"][f]["tier_applied"] for i in tiered)]
    gem = TIERS["gemini/gemini-3.1-pro-preview"]
    extra_tokens = max(m["context"] for m in models if m["id"].startswith("gemini/")) - 1_000_000
    return {
        "priciest_full": {"id": hi, "cost": round(full(hi), 4)},
        "cheapest_full": {"id": lo, "cost": round(full(lo), 4)},
        "spread_full": round(full(hi) / full(lo), 1),
        "priciest_full_long_context": {"id": lhi, "cost": round(full(lhi), 4)},
        "cheapest_full_long_context": {"id": llo, "cost": round(full(llo), 4)},
        "spread_full_long_context": round(full(lhi) / full(llo), 1),
        "cost_200k_min": {"ids": [i for i in ids if abs(raw[i]["c200"] - c200_min) < 1e-12],
                          "cost": round(c200_min, 4)},
        "cost_200k_max": {"ids": [i for i in ids if abs(raw[i]["c200"] - c200_max) < 1e-12],
                          "cost": round(c200_max, 4)},
        "spread_200k": round(c200_max / c200_min, 1),
        "tiered_models": tiered,
        "tier_free_fractions": free,
        "long_context_ids": long_ids,
        "long_context_steps": steps,
        "fits_but_pricey": [i for i in long_ids if by_id[i]["fits_but_pricey"]],
        "fits_but_cheap": [i for i in long_ids if not by_id[i]["fits_but_pricey"]],
        "gemini_extra_tokens": extra_tokens,
        "gemini_extra_cost_at_tier": round(extra_tokens * gem["in_above"] / 1e6, 4),
    }


# ---------------------------------------------------------------- display strings
def usd4(x):
    return "${:.4f}".format(x)


def usd2(x):
    return "${:.2f}".format(x)


def mult1(x):
    return "{:.1f}×".format(x)


def tok(n):
    return "{:,}".format(n)


def build_display(pricing, models, summary):
    """The final string of every number the article's prose may use."""
    by_id = {m["id"]: m for m in models}
    s = summary
    d = {
        "pricing_last_synced": pricing["last_synced"],
        "tiers_read_on": TIERS_READ_ON,
        "n_models": str(len(models)),
        "n_long_context": str(len(s["long_context_ids"])),
        "n_tiered": str(len(s["tiered_models"])),
        "n_fits_but_pricey": str(len(s["fits_but_pricey"])),
        "n_fits_but_cheap": str(len(s["fits_but_cheap"])),
        "long_context_min": tok(LONG_CONTEXT_MIN),
        "common_tokens": tok(COMMON_TOKENS),
        "pricey_threshold": usd2(PRICEY_THRESHOLD_USD),
        "budget": "${:,.0f}".format(BUDGET_USD),
        "priciest_full": usd4(s["priciest_full"]["cost"]),
        "cheapest_full": usd4(s["cheapest_full"]["cost"]),
        "spread_full": mult1(s["spread_full"]),
        "priciest_full_long": usd4(s["priciest_full_long_context"]["cost"]),
        "cheapest_full_long": usd4(s["cheapest_full_long_context"]["cost"]),
        "spread_full_long": mult1(s["spread_full_long_context"]),
        "cost_200k_min": usd4(s["cost_200k_min"]["cost"]),
        "cost_200k_max": usd4(s["cost_200k_max"]["cost"]),
        "spread_200k": mult1(s["spread_200k"]),
        "gemini_extra_tokens": tok(s["gemini_extra_tokens"]),
        "gemini_extra_cost": usd4(s["gemini_extra_cost_at_tier"]),
        "step1_low": usd4(s["long_context_steps"][0]["low_cost"]),
        "step1_high": usd4(s["long_context_steps"][0]["high_cost"]),
        "step1_ratio": mult1(s["long_context_steps"][0]["ratio"]),
        "step2_low": usd4(s["long_context_steps"][1]["low_cost"]),
        "step2_high": usd4(s["long_context_steps"][1]["high_cost"]),
        "step2_ratio": mult1(s["long_context_steps"][1]["ratio"]),
        "model": {}, "tier": {},
    }
    for m in models:
        f = m["fill"]
        d["model"][m["id"]] = {
            "context": tok(m["context"]), "in": usd2(m["in"]), "out": usd2(m["out"]),
            "tokens_full": tok(f["full"]["tokens"]),
            "tokens_half": tok(f["half"]["tokens"]),
            "tokens_tenth": tok(f["tenth"]["tokens"]),
            "full": usd4(f["full"]["cost_applicable"]),
            "half": usd4(f["half"]["cost_applicable"]),
            "tenth": usd4(f["tenth"]["cost_applicable"]),
            "full_flat": usd4(f["full"]["cost_flat"]),
            "half_flat": usd4(f["half"]["cost_flat"]),
            "tenth_flat": usd4(f["tenth"]["cost_flat"]),
            "cost_200k": usd4(m["cost_200k"]),
            "requests_per_100usd_full": "{:.1f}".format(m["requests_per_100usd_full"]),
            "full_over_tenth": mult1(m["full_over_tenth"]),
        }
        t = m["tier"]
        if t:
            d["tier"][m["id"]] = {
                "above": tok(t["above"]),
                "in_base": usd2(t["in_base"]), "in_above": usd2(t["in_above"]),
                "out_base": usd2(t["out_base"]), "out_above": usd2(t["out_above"]),
                "in_mult": mult1(t["in_above"] / t["in_base"]),
                "out_mult": mult1(t["out_above"] / t["out_base"]),
            }
    return d


# ---------------------------------------------------------------- checks
def run_checks(pricing, models, summary):
    chat = pricing["chat"]
    by_pid = {m["id"]: m for m in chat}
    out = {}
    out["1_pricing_has_13_valid_models"] = (
        len(chat) == 13 and all(m["in"] > 0 and m["out"] > 0 and m["context"] > 0 for m in chat))
    out["2_tiers_base_rates_match_pricing"] = all(
        i in by_pid and by_pid[i]["in"] == t["in_base"] and by_pid[i]["out"] == t["out_base"]
        for i, t in TIERS.items())
    ok3 = True
    for m in models:
        src = by_pid[m["id"]]
        for name, frac in FILL_FRACTIONS.items():
            f = m["fill"][name]
            exact = float(round(Decimal(f["tokens"]) * Decimal(str(src["in"])) / Decimal(10 ** 6), 4))
            ok3 = ok3 and (f["tokens"] == round(src["context"] * frac)
                           and abs(f["tokens"] * src["in"] / 1e6 - f["cost_flat"]) < 1e-4
                           and f["cost_flat"] == exact)
    out["3_tokens_and_flat_cost_consistent"] = ok3
    ok4 = True
    for m in models:
        a = [m["fill"][n]["tier_applied"] for n in ("full", "half", "tenth")]
        ok4 = ok4 and (a == [True, True, False] if m["tier"] else a == [False, False, False])
    out["4_tier_applied_pattern"] = ok4
    st = summary["long_context_steps"][0]
    out["5_pricey_threshold_in_largest_step"] = st["low_cost"] < PRICEY_THRESHOLD_USD <= st["high_cost"]
    out["6_long_context_partition_complete"] = (
        len(summary["fits_but_pricey"]) + len(summary["fits_but_cheap"]) == len(summary["long_context_ids"]) > 0)
    out["7_common_tokens_is_smallest_max_input"] = COMMON_TOKENS == min(m["context"] for m in chat)
    return out


# ---------------------------------------------------------------- html output
def table1_rows(models):
    rows = []
    for m in models:
        f, t = m["fill"], m["tier"]
        mark = " †" if t else ""
        rate = "{:.2f}".format(m["in"])
        if t:
            rate += " ({:.2f} above {}K)".format(t["in_above"], t["above"] // 1000)
        rows.append("<tr><td>{}{}</td><td>{:,}</td><td>{}</td><td>{:.4f}</td><td>{:.4f}</td><td>{:.4f}</td></tr>".format(
            m["name"], mark, m["context"], rate,
            f["full"]["cost_applicable"], f["half"]["cost_applicable"], f["tenth"]["cost_applicable"]))
    return "\n".join(rows)


def table2_rows(models):
    rows = []
    for m in models:
        if not m["long_context"]:
            continue
        rows.append("<tr><td>{}</td><td>{:.4f}</td><td>{:.1f}</td><td>{:.4f}</td></tr>".format(
            m["name"], m["fill"]["full"]["cost_applicable"], m["requests_per_100usd_full"], m["cost_200k"]))
    return "\n".join(rows)


def figure_svg(models):
    x0, scale, y0, pitch = 200, 40.0, 56, 24     # 40 px per USD, axis 0..10
    axis_max = 10
    n = len(models)
    h = y0 + n * pitch + 54
    parts = ['<svg viewBox="0 0 720 {}" width="100%" role="img" font-family="Noto Sans,Segoe UI,sans-serif">'.format(h),
             "<title>Cost of one request that fills the input window, 13 models</title>",
             "<desc>Horizontal bars: the price in US dollars of one request whose prompt is the model's full maximum "
             "input size, input side only. A lighter segment marks the part added by the higher rate that four models "
             "charge above 200K or 272K tokens. Values are in Table 1.</desc>"]
    for g in range(0, axis_max + 1, 2):
        gx = x0 + g * scale
        parts.append('<line x1="{0}" y1="{1}" x2="{0}" y2="{2}" stroke="#e2e2e2"/>'.format(
            gx, y0 - 6, y0 + n * pitch))
        parts.append('<text x="{}" y="{}" font-size="11" text-anchor="middle" fill="#6b7280">${}</text>'.format(
            gx, y0 + n * pitch + 16, g))
    lx = 200
    for vendor, color in VENDOR_COLORS.items():
        parts.append('<rect x="{}" y="10" width="12" height="12" fill="{}"/>'.format(lx, color))
        parts.append('<text x="{}" y="20" font-size="12" fill="#1f2937">{}</text>'.format(lx + 17, vendor))
        lx += 110
    parts.append('<rect x="200" y="28" width="12" height="12" fill="#6b7280" fill-opacity="0.35"/>')
    parts.append('<text x="217" y="38" font-size="12" fill="#1f2937">lighter segment = added by the higher rate above 200K / 272K tokens</text>')
    for i, m in enumerate(models):
        y = y0 + i * pitch
        color = VENDOR_COLORS[m["vendor"]]
        appl, flat = m["fill"]["full"]["cost_applicable"], m["fill"]["full"]["cost_flat"]
        w_appl, w_flat = appl * scale, flat * scale
        label = m["name"] + (" †" if m["tier"] else "")
        parts.append('<text x="190" y="{}" font-size="12" text-anchor="end" fill="#1f2937">{}</text>'.format(y + 14, label))
        parts.append('<rect x="{}" y="{}" width="{:.1f}" height="16" fill="{}"/>'.format(x0, y + 3, w_flat, color))
        if m["tier"]:
            parts.append('<rect x="{:.1f}" y="{}" width="{:.1f}" height="16" fill="{}" fill-opacity="0.35"/>'.format(
                x0 + w_flat, y + 3, w_appl - w_flat, color))
        parts.append('<text x="{:.1f}" y="{}" font-size="11" fill="#1f2937">${:.4f}</text>'.format(
            x0 + w_appl + 6, y + 15, appl))
    parts.append('<text x="200" y="{}" font-size="11" fill="#6b7280">USD for one request that fills the whole input window (input side only)</text>'.format(
        y0 + n * pitch + 34))
    parts.append("</svg>")
    return "".join(parts)


# ---------------------------------------------------------------- main
def main():
    if hasattr(sys.stdout, "reconfigure"):
        sys.stdout.reconfigure(encoding="utf-8", newline="\n")  # table/figure output has a dagger; LF everywhere
    pricing = json.loads(PRICING_PATH.read_text(encoding="utf-8"))
    models, raw = build_models(pricing)
    summary = build_summary(models, raw)
    checks = run_checks(pricing, models, summary)

    print("Pricing source: {} (last synced {})".format(pricing["source_name"], pricing["last_synced"]))
    for name, ok in checks.items():
        print("  check {}: {}".format(name, "True" if ok else "FALSE"))
    if not all(checks.values()):
        print("Some checks failed: nothing written.")
        raise SystemExit(1)

    results = {
        "pricing_source": pricing["source_name"],
        "pricing_last_synced": pricing["last_synced"],
        "tiers_read_on": TIERS_READ_ON,
        "assumptions": {
            "fill_fractions": FILL_FRACTIONS,
            "common_tokens": COMMON_TOKENS,
            "budget_usd": BUDGET_USD,
            "long_context_min": LONG_CONTEXT_MIN,
            "pricey_threshold_usd": PRICEY_THRESHOLD_USD,
            "context_meaning": "max_input_tokens in the LiteLLM dataset: an input limit, not the output limit",
            "tier_rule": "once prompt tokens are strictly greater than the threshold, the whole prompt "
                         "is billed at the higher input rate and the output rate switches too "
                         "(as LiteLLM's cost calculator applies it; the vendor's bill is the authority)",
            "ratio_rule": "ratios and fills-per-$100 are computed from unrounded costs",
            "tiers_dataset_url": TIERS_DATASET_URL,
            "tiers_dataset_md5": TIERS_DATASET_MD5,
            "tiers_calc_commit": TIERS_CALC_COMMIT,
        },
        "models": models,
        "summary": summary,
        "display": build_display(pricing, models, summary),
        "checks": checks,
    }

    print()
    print("{:<30}{:>11}{:>12}{:>12}{:>12}{:>12}".format("Model", "Max input", "Full $", "Half $", "Tenth $", "200K $"))
    for m in models:
        f = m["fill"]
        print("{:<30}{:>11,}{:>12.4f}{:>12.4f}{:>12.4f}{:>12.4f}".format(
            m["name"] + (" *" if m["tier"] else ""), m["context"], f["full"]["cost_applicable"],
            f["half"]["cost_applicable"], f["tenth"]["cost_applicable"], m["cost_200k"]))
    print("(* = long-context rate applied above the threshold)")
    print()
    print("fits_but_pricey ({}): {}".format(len(summary["fits_but_pricey"]), ", ".join(summary["fits_but_pricey"])))
    print("fits_but_cheap ({}): {}".format(len(summary["fits_but_cheap"]), ", ".join(summary["fits_but_cheap"])))

    with open(RESULTS_PATH, "w", encoding="utf-8", newline="\n") as fh:  # LF on every OS: same bytes everywhere
        fh.write(json.dumps(results, ensure_ascii=False, indent=2) + "\n")
    print()
    print("Wrote {}".format(RESULTS_PATH))

    if "--html" in sys.argv[1:]:
        print()
        print("### table1_tbody")
        print(table1_rows(models))
        print("### table2_tbody")
        print(table2_rows(models))
        print("### figure1_svg")
        print(figure_svg(models))
        print("### end")


if __name__ == "__main__":
    main()
