#!/usr/bin/env python3
"""First-party global USD estimate; 2026-10-15 rates. No API calls."""
import argparse
from decimal import Decimal as D


def estimate(ordinary=0, output=0, write5=0, write1=0, read=0, batch=False):
    counts = (ordinary, output, write5, write1, read)
    if any(type(n) is not int or n < 0 for n in counts):
        raise ValueError('Token counts must be nonnegative integers')
    if type(batch) is not bool:
        raise ValueError('batch must be a boolean')
    prompt = ordinary + write5 + write1 + read
    high = prompt > 100000
    rates = ('0.50', '2.50', '0.625', '1.00', '0.05') if high else ('0.10', '0.50', '0.125', '0.20', '0.01')
    cost = sum(D(n) * D(r) for n, r in zip(counts, rates)) / D(1000000)
    if batch:
        cost *= D('0.5')
    return {'prompt_tokens': prompt, 'tier': 'higher' if high else 'lower', 'usd': format(cost, 'f')}


def self_test():
    assert D(estimate(99999, 1000)['usd']) == D('0.0104999')
    assert D(estimate(100000, 1000)['usd']) == D('0.0105')
    assert D(estimate(100001, 1000)['usd']) == D('0.0525005')
    assert estimate(99000, 2000)['tier'] == 'lower'
    assert D(estimate(10001, 1000, read=90000)['usd']) == D('0.0120005')
    assert D(estimate(10001, 1000, read=90000, batch=True)['usd']) == D('0.00600025')
    assert D(estimate(write5=100000)['usd']) == D('0.0125')
    assert D(estimate(write1=100000)['usd']) == D('0.020')
    assert D(estimate(write5=100001)['usd']) == D('0.062500625')
    assert D(estimate(write1=100001)['usd']) == D('0.100001')
    assert estimate(write5=50000, write1=50001)['tier'] == 'higher'
    for bad in (-1, 1.2, True, '1'):
        try:
            estimate(bad)
        except ValueError:
            pass
        else:
            raise AssertionError('Accepted invalid token count')
    print('Self-test passed: published formula, not production billing verification.')


if __name__ == '__main__':
    p = argparse.ArgumentParser(description=__doc__)
    for flag in ('input', 'output', 'cache-write-5m', 'cache-write-1h', 'cache-read'):
        p.add_argument('--' + flag, type=int, default=0)
    p.add_argument('--batch', action='store_true')
    p.add_argument('--self-test', action='store_true')
    a = p.parse_args()
    if a.self_test:
        self_test()
    else:
        try:
            result = estimate(a.input, a.output, a.cache_write_5m, a.cache_write_1h, a.cache_read, a.batch)
        except ValueError as e:
            p.error(str(e))
        import json
        print(json.dumps(result, indent=2))