"""what-one-coding-task-costs-tokens-dollars-minutes.py — price a coding-agent run from token counts.

What it does:
  1. Holds a dated price table (USD per million tokens) copied from the vendors' pricing pages
     on 2026-09-05. Cache-read tokens are priced at 0.1x base input and 5-minute cache writes at
     1.25x base input for the Claude rows, as the Claude pricing page states; the OpenAI rows carry
     the "cached input" price the OpenAI pricing page lists.
  2. cost(model, input, cache_write, cache_read, output) -> dollars.
  3. estimate_tokens(text) -> int using the chars/4 heuristic (the pricing FAQ's "1 token is
     approximately 4 characters"). It is NOT a tokenizer; expect +-30% on code.
  4. Emits CSV rows in the dataset's column order, and can self-check against two published
     worked examples (--selftest).

Inputs / run:
  python what-one-coding-task-costs-tokens-dollars-minutes.py --selftest
  python what-one-coding-task-costs-tokens-dollars-minutes.py --model claude-opus-5 --in 50000 --out 15000
  python what-one-coding-task-costs-tokens-dollars-minutes.py --estimate path/to/prompt.txt --model claude-sonnet-5
Needs: Python 3.10+ (tested 3.13.12 on Windows). Standard library only.
"""
from __future__ import annotations

import argparse
import csv
import sys

PRICES_READ_AT = "2026-09-05"
# model: (base input $/MTok, 5m cache write $/MTok, cache read $/MTok, output $/MTok, source)
PRICES = {
    "claude-opus-5": (5.00, 6.25, 0.50, 25.00, "https://platform.claude.com/docs/en/about-claude/pricing"),
    "claude-sonnet-5": (2.00, 2.50, 0.20, 10.00, "https://platform.claude.com/docs/en/about-claude/pricing"),
    "claude-sonnet-4-6": (3.00, 3.75, 0.30, 15.00, "https://platform.claude.com/docs/en/about-claude/pricing"),
    "claude-haiku-4-5": (1.00, 1.25, 0.10, 5.00, "https://platform.claude.com/docs/en/about-claude/pricing"),
    # OpenAI rows: the pricing page lists input / cached input / output; no separate cache-write price,
    # so cache_write is charged at the base input price here.
    "gpt-5.5": (5.00, 5.00, 0.50, 30.00, "https://developers.openai.com/api/docs/pricing"),
    "gpt-5.3-codex": (1.75, 1.75, 0.175, 14.00, "https://developers.openai.com/api/docs/pricing"),
    "gpt-5-mini": (0.25, 0.25, 0.025, 2.00, "https://developers.openai.com/api/docs/pricing"),
}

COLUMNS = ["row_id", "source", "model", "input_tokens", "cache_write_tokens", "cache_read_tokens",
           "output_tokens", "price_in_per_mtok", "price_out_per_mtok", "cost_usd", "api_minutes",
           "approach", "note"]


def cost(model: str, inp: int, cache_write: int, cache_read: int, out: int) -> float:
    p_in, p_cw, p_cr, p_out, _ = PRICES[model]
    return (inp * p_in + cache_write * p_cw + cache_read * p_cr + out * p_out) / 1_000_000


def estimate_tokens(text: str) -> int:
    """chars/4 heuristic. Not a tokenizer."""
    return max(1, round(len(text) / 4))


def row(row_id: str, source: str, model: str, inp: int, cw: int, cr: int, out: int,
        api_minutes: str, approach: str, note: str) -> dict:
    p = PRICES[model]
    return {"row_id": row_id, "source": source, "model": model, "input_tokens": inp,
            "cache_write_tokens": cw, "cache_read_tokens": cr, "output_tokens": out,
            "price_in_per_mtok": p[0], "price_out_per_mtok": p[3],
            "cost_usd": f"{cost(model, inp, cw, cr, out):.4f}", "api_minutes": api_minutes,
            "approach": approach, "note": note}


def selftest() -> int:
    """Reproduce the published worked examples; return 0 if all match."""
    checks = [
        # Claude pricing page, Managed Agents worked example (tokens only, session runtime excluded):
        ("claude-opus-5", 50_000, 0, 0, 15_000, 0.625, "pricing page: 50k in + 15k out on Opus 5 = $0.25 + $0.375"),
        ("claude-opus-5", 10_000, 0, 40_000, 15_000, 0.445, "pricing page: 10k in + 40k cache read + 15k out = $0.05 + $0.02 + $0.375"),
        # Claude Code costs page /usage block: 1.2k input, 5.3k output, 940.0k cache read, 50.0k cache write ($0.55)
        ("claude-sonnet-4-6", 1_200, 50_000, 940_000, 5_300, 0.55, "costs page /usage block (rounded display values)"),
    ]
    bad = 0
    for model, i, cw, cr, o, expected, label in checks:
        got = cost(model, i, cw, cr, o)
        ok = abs(got - expected) < 0.01
        bad += not ok
        print(f"{'ok  ' if ok else 'FAIL'} {label}: computed ${got:.4f} expected ${expected:.3f}")
    return bad


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    ap.add_argument("--selftest", action="store_true")
    ap.add_argument("--model", choices=sorted(PRICES), default="claude-sonnet-5")
    ap.add_argument("--in", dest="inp", type=int, default=0)
    ap.add_argument("--cache-write", type=int, default=0)
    ap.add_argument("--cache-read", type=int, default=0)
    ap.add_argument("--out", type=int, default=0)
    ap.add_argument("--estimate", help="text file: estimate its input tokens with chars/4")
    ap.add_argument("--csv", action="store_true", help="print one CSV row instead of a sentence")
    a = ap.parse_args()
    if a.selftest:
        sys.exit(selftest())
    inp = a.inp
    if a.estimate:
        text = open(a.estimate, encoding="utf-8").read()
        inp += estimate_tokens(text)
        print(f"{a.estimate}: {len(text)} chars -> ~{estimate_tokens(text)} tokens (chars/4 heuristic)")
    usd = cost(a.model, inp, a.cache_write, a.cache_read, a.out)
    if a.csv:
        w = csv.DictWriter(sys.stdout, fieldnames=COLUMNS, lineterminator="\n")
        w.writeheader()
        w.writerow(row("manual", "command line", a.model, inp, a.cache_write, a.cache_read, a.out, "", "", ""))
    else:
        print(f"{a.model}: {inp} in, {a.cache_write} cache write, {a.cache_read} cache read, {a.out} out"
              f" -> ${usd:.4f} at prices read {PRICES_READ_AT}")


if __name__ == "__main__":
    main()