#!/usr/bin/env python3
"""Offline synthetic token-cost example; not an API adapter or invoice calculator."""
import copy
from decimal import Decimal, InvalidOperation, localcontext
import json
import sys
import unittest

COUNTS = {"input", "cached_input", "cache_write", "output", "reasoning"}
RATES = {"uncached_input", "cached_input", "cache_write", "output"}
EXAMPLE = {
    "counts": {"input": 10000, "cached_input": 6000, "cache_write": 2000,
               "output": 1000, "reasoning": 700},
    "usd_per_million": {"uncached_input": "2", "cached_input": "0.2",
                        "cache_write": "2.5", "output": "8"}
}


def exact_keys(value, expected, label):
    if not isinstance(value, dict) or set(value) != expected:
        raise ValueError(label + " must contain exactly " + ", ".join(sorted(expected)))


def estimate(record):
    exact_keys(record, {"counts", "usd_per_million"}, "record")
    c, raw = record["counts"], record["usd_per_million"]
    exact_keys(c, COUNTS, "counts")
    exact_keys(raw, RATES, "usd_per_million")
    for name, value in c.items():
        if type(value) is not int or not 0 <= value <= 10**12:
            raise ValueError(name + " must be an integer between 0 and 10^12")
    if c["cached_input"] + c["cache_write"] > c["input"]:
        raise ValueError("cached input plus cache writes exceeds input")
    if c["reasoning"] > c["output"]:
        raise ValueError("reasoning exceeds output")
    rates = {}
    for name, value in raw.items():
        if not isinstance(value, str) or not value or len(value) > 30:
            raise ValueError(name + " rate must be a short decimal string")
        try:
            rate = Decimal(value)
        except InvalidOperation as exc:
            raise ValueError(name + " rate is invalid") from exc
        if not rate.is_finite() or not 0 <= rate <= 1000000:
            raise ValueError(name + " rate must be finite and nonnegative")
        exponent = rate.as_tuple().exponent
        if not isinstance(exponent, int) or exponent < -12:
            raise ValueError(name + " rate supports at most 12 decimal places")
        rates[name] = rate
    buckets = {
        "uncached_input": c["input"] - c["cached_input"] - c["cache_write"],
        "cached_input": c["cached_input"], "cache_write": c["cache_write"],
        "output": c["output"]
    }
    with localcontext() as ctx:
        ctx.prec = 60
        costs = {k: Decimal(n) * rates[k] / Decimal(1000000)
                 for k, n in buckets.items()}
        total = sum(costs.values(), Decimal(0))
    return {"billable_buckets": buckets,
            "usd_by_bucket": {k: format(v, "f") for k, v in costs.items()},
            "estimated_token_usd": format(total, "f"),
            "reasoning_subset_of_output": c["reasoning"]}


def unique_object(pairs):
    result = {}
    for key, value in pairs:
        if key in result:
            raise ValueError("duplicate JSON key: " + key)
        result[key] = value
    return result


class Tests(unittest.TestCase):
    def test_example(self):
        self.assertEqual(estimate(EXAMPLE)["estimated_token_usd"], "0.0182")

    def test_no_cache(self):
        r = copy.deepcopy(EXAMPLE)
        r["counts"].update(cached_input=0, cache_write=0)
        self.assertEqual(estimate(r)["estimated_token_usd"], "0.028")

    def test_reasoning_is_subset(self):
        r = copy.deepcopy(EXAMPLE)
        r["counts"]["reasoning"] = 0
        self.assertEqual(estimate(r)["estimated_token_usd"], "0.0182")

    def test_zero_counts(self):
        r = copy.deepcopy(EXAMPLE)
        r["counts"] = dict.fromkeys(COUNTS, 0)
        self.assertEqual(Decimal(estimate(r)["estimated_token_usd"]), 0)

    def test_bad_count_types_and_bounds(self):
        for value in (True, -1, 1.5, "1", None, 10**12 + 1):
            with self.subTest(value=value):
                r = copy.deepcopy(EXAMPLE)
                r["counts"]["input"] = value
                with self.assertRaises(ValueError):
                    estimate(r)

    def test_overlapping_input(self):
        r = copy.deepcopy(EXAMPLE)
        r["counts"]["cache_write"] = 5000
        with self.assertRaises(ValueError):
            estimate(r)

    def test_reasoning_exceeds_output(self):
        r = copy.deepcopy(EXAMPLE)
        r["counts"]["reasoning"] = 1001
        with self.assertRaises(ValueError):
            estimate(r)

    def test_missing_usage(self):
        r = copy.deepcopy(EXAMPLE)
        del r["counts"]["cached_input"]
        with self.assertRaises(ValueError):
            estimate(r)

    def test_unknown_category(self):
        r = copy.deepcopy(EXAMPLE)
        r["counts"]["audio"] = 50
        with self.assertRaises(ValueError):
            estimate(r)

    def test_bad_rates(self):
        for value in (2.0, "NaN", "Infinity", "-1", "bad", "0.0000000000001"):
            with self.subTest(value=value):
                r = copy.deepcopy(EXAMPLE)
                r["usd_per_million"]["output"] = value
                with self.assertRaises(ValueError):
                    estimate(r)

    def test_small_charge_retained(self):
        r = copy.deepcopy(EXAMPLE)
        r["counts"] = dict.fromkeys(COUNTS, 0)
        r["counts"]["input"] = 1
        self.assertEqual(estimate(r)["estimated_token_usd"], "0.000002")

    def test_duplicate_json_key(self):
        with self.assertRaises(ValueError):
            json.loads('{"counts": {}, "counts": {}}', object_pairs_hook=unique_object)


if __name__ == "__main__":
    if len(sys.argv) == 1:
        print(json.dumps(estimate(EXAMPLE), indent=2))
        result = unittest.TextTestRunner(verbosity=2).run(
            unittest.defaultTestLoader.loadTestsFromTestCase(Tests))
        sys.exit(0 if result.wasSuccessful() else 1)
    if len(sys.argv) != 2:
        sys.exit("usage: python3 token-cost-partition.py [normalized-record.json]")
    try:
        with open(sys.argv[1], encoding="utf-8") as source:
            record = json.load(source, object_pairs_hook=unique_object)
        print(json.dumps(estimate(record), indent=2))
    except (ValueError, OSError, InvalidOperation) as exc:
        print("ERROR: " + str(exc), file=sys.stderr)
        sys.exit(1)