#!/usr/bin/env python3 """Calculate API-compatible user_unit_cost totals from exact JSONL amounts.""" from __future__ import annotations import argparse import decimal import hashlib import json from collections import OrderedDict from pathlib import Path from typing import Any, Iterable BASE_GROUP_FIELDS = ( "tenant_id", "unit", "labels", "consumed_source", "invoice_id", "internal", ) DEFAULT_AMOUNT_FIELDS = ("cost_price", "raw_price") CREDITS_CONSUMED_SOURCE = "credits" def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser( description=( "Group full-precision user_unit_cost amounts, truncate each group " "toward zero, and aggregate the truncated groups." ) ) parser.add_argument("--input", type=Path, required=True, help="Input JSONL path.") parser.add_argument("--output", type=Path, required=True, help="Result JSON path.") parser.add_argument("--precision", type=int, default=2) parser.add_argument( "--amount-fields", default=",".join(DEFAULT_AMOUNT_FIELDS), help="Comma-separated exact decimal fields (default: cost_price,raw_price).", ) parser.add_argument("--ignore-invoice-id", action="store_true") labels_group = parser.add_mutually_exclusive_group() labels_group.add_argument( "--include-labels", dest="ignore_labels", action="store_false", help="Use standard grouping with labels (the default without a decision file).", ) labels_group.add_argument( "--ignore-labels", dest="ignore_labels", action="store_true", help="Explicitly use grouping without labels.", ) labels_group.add_argument( "--grouping-decision", type=Path, help="Decision JSON produced by resolve_user_unit_cost_grouping.py.", ) parser.set_defaults(ignore_labels=False, grouping_decision=None) parser.add_argument( "--include-billing-day", action="store_true", help="Add billing_day to the group key for the reviewed June 2026 freeze only.", ) parser.add_argument( "--rollup-field", action="append", default=[], metavar="NAME=PATH", help=( "Optional reporting rollup, repeatable (for example " "sku=unit or source=consumed_source). It is applied only after " "API-compatible per-group truncation." ), ) args = parser.parse_args() if not 0 <= args.precision <= 30: parser.error("--precision must be between 0 and 30.") args.amount_fields = tuple( field.strip() for field in args.amount_fields.split(",") if field.strip() ) if not args.amount_fields: parser.error("--amount-fields must contain at least one field.") try: args.rollup_fields = parse_rollup_fields(args.rollup_field) except ValueError as exc: parser.error(str(exc)) return args def parse_rollup_fields(values: list[str]) -> tuple[tuple[str, str], ...]: result: list[tuple[str, str]] = [] aliases: set[str] = set() for value in values: if "=" not in value: raise ValueError("--rollup-field must use NAME=PATH syntax.") alias, path = (part.strip() for part in value.split("=", 1)) if not alias or not path or any(not part for part in path.split(".")): raise ValueError("--rollup-field requires a non-empty NAME and PATH.") if alias in aliases: raise ValueError(f"Duplicate rollup field name: {alias}.") aliases.add(alias) result.append((alias, path)) return tuple(result) def normalize_labels(value: Any) -> Any: if not isinstance(value, str): return value try: parsed = json.loads(value) except json.JSONDecodeError: return value return parsed def canonical_value(field: str, value: Any) -> tuple[str, Any]: normalized = normalize_labels(value) if field == "labels" else value canonical = json.dumps( normalized, ensure_ascii=False, sort_keys=True, separators=(",", ":") ) return canonical, normalized def resolve_path(row: dict[str, Any], path: str, *, line_number: int) -> Any: current: Any = row for index, part in enumerate(path.split(".")): if index == 1 and path.split(".", 1)[0] == "labels": current = normalize_labels(current) if not isinstance(current, dict) or part not in current: raise ValueError(f"Line {line_number}: rollup path {path!r} is missing.") current = current[part] return current def parse_decimal(value: Any, *, field: str, line_number: int) -> decimal.Decimal: if isinstance(value, bool) or value is None: raise ValueError(f"Line {line_number}: {field} must be an exact decimal value.") try: result = decimal.Decimal(str(value)) except decimal.InvalidOperation as exc: raise ValueError( f"Line {line_number}: {field} is not an exact decimal value." ) from exc if not result.is_finite(): raise ValueError(f"Line {line_number}: {field} must be finite.") return result def truncate_toward_zero(value: decimal.Decimal, precision: int) -> decimal.Decimal: quantum = decimal.Decimal(1).scaleb(-precision) truncated = value.quantize(quantum, rounding=decimal.ROUND_DOWN) return abs(truncated) if truncated.is_zero() else truncated def fixed_decimal(value: decimal.Decimal, precision: int) -> str: return format(value, f".{precision}f") def exact_decimal(value: decimal.Decimal) -> str: return format(value, "f") def load_jsonl(path: Path) -> Iterable[tuple[int, dict[str, Any]]]: with path.open(encoding="utf-8") as handle: for line_number, line in enumerate(handle, start=1): if not line.strip(): continue try: value = json.loads(line) except json.JSONDecodeError as exc: raise ValueError(f"Line {line_number}: invalid JSON.") from exc if not isinstance(value, dict): raise ValueError(f"Line {line_number}: expected a JSON object.") yield line_number, value def calculate( rows: Iterable[tuple[int, dict[str, Any]]], *, amount_fields: tuple[str, ...], precision: int, ignore_invoice_id: bool = False, ignore_labels: bool = False, include_billing_day: bool = False, rollup_fields: tuple[tuple[str, str], ...] = (), ) -> dict[str, Any]: group_fields = list(BASE_GROUP_FIELDS) if ignore_invoice_id: group_fields.remove("invoice_id") if ignore_labels: group_fields.remove("labels") if include_billing_day: group_fields.append("billing_day") groups: dict[tuple[str, ...], dict[str, Any]] = {} input_row_count = 0 with decimal.localcontext() as context: context.prec = 100 for line_number, row in rows: input_row_count += 1 missing = [field for field in (*group_fields, *amount_fields) if field not in row] if missing: raise ValueError( f"Line {line_number}: missing required fields: {', '.join(missing)}." ) key_parts: list[str] = [] dimensions: OrderedDict[str, Any] = OrderedDict() for field in group_fields: canonical, normalized = canonical_value(field, row[field]) key_parts.append(canonical) dimensions[field] = normalized key = tuple(key_parts) rollup_dimensions: OrderedDict[str, Any] = OrderedDict() rollup_key_parts: list[str] = [] for alias, path in rollup_fields: value = resolve_path(row, path, line_number=line_number) canonical, normalized = canonical_value(alias, value) rollup_key_parts.append(canonical) rollup_dimensions[alias] = normalized group = groups.setdefault( key, { "dimensions": dimensions, "rollup_dimensions": rollup_dimensions, "rollup_key": tuple(rollup_key_parts), "full_precision": { field: decimal.Decimal(0) for field in amount_fields }, }, ) if group["rollup_key"] != tuple(rollup_key_parts): raise ValueError( f"Line {line_number}: one API precision group maps to multiple " "rollup values; preserve the upstream grouping dimension instead " "of splitting the group during calculation." ) for field in amount_fields: group["full_precision"][field] += parse_decimal( row[field], field=field, line_number=line_number ) totals = {field: decimal.Decimal(0) for field in amount_fields} full_precision_totals = {field: decimal.Decimal(0) for field in amount_fields} has_billing_amounts = {"cost_price", "raw_price"}.issubset(amount_fields) billing_overview = { "runningTotal": decimal.Decimal(0), "discounts": decimal.Decimal(0), "credits": decimal.Decimal(0), "totalCost": decimal.Decimal(0), } rollup_groups: dict[tuple[str, ...], dict[str, Any]] = {} output_groups = [] for key in sorted(groups): group = groups[key] truncated = { field: truncate_toward_zero(group["full_precision"][field], precision) for field in amount_fields } for field in amount_fields: totals[field] += truncated[field] full_precision_totals[field] += group["full_precision"][field] if has_billing_amounts: raw_price = truncated["raw_price"] cost_price = truncated["cost_price"] billing_overview["runningTotal"] += raw_price billing_overview["discounts"] += raw_price - cost_price if group["dimensions"]["consumed_source"] == CREDITS_CONSUMED_SOURCE: billing_overview["credits"] += cost_price else: billing_overview["totalCost"] += cost_price output_groups.append( { "dimensions": group["dimensions"], **( {"rollup_dimensions": group["rollup_dimensions"]} if rollup_fields else {} ), "full_precision": { field: exact_decimal(group["full_precision"][field]) for field in amount_fields }, "truncated": { field: fixed_decimal(truncated[field], precision) for field in amount_fields }, } ) if rollup_fields: rollup = rollup_groups.setdefault( group["rollup_key"], { "dimensions": group["rollup_dimensions"], "totals": { field: decimal.Decimal(0) for field in amount_fields }, }, ) for field in amount_fields: rollup["totals"][field] += truncated[field] output_rollups = [ { "dimensions": rollup_groups[key]["dimensions"], "totals": { field: fixed_decimal(rollup_groups[key]["totals"][field], precision) for field in amount_fields }, } for key in sorted(rollup_groups) ] if ignore_labels and ignore_invoice_id: grouping_mode = "ignore_labels_and_invoice_id" elif ignore_labels: grouping_mode = "ignore_labels" elif ignore_invoice_id: grouping_mode = "ignore_invoice_id" else: grouping_mode = "standard" result = { "format": "manual-invoice-user-unit-cost/v2", "algorithm": "group -> sum full precision -> truncate each group toward zero -> aggregate", "precision": precision, "grouping_mode": grouping_mode, "group_fields": group_fields, "amount_fields": list(amount_fields), "rollup_fields": [ {"name": alias, "path": path} for alias, path in rollup_fields ], "input_row_count": input_row_count, "group_count": len(output_groups), "full_precision_totals": { field: exact_decimal(full_precision_totals[field]) for field in amount_fields }, "totals": { field: fixed_decimal(totals[field], precision) for field in amount_fields }, "rollups": output_rollups, "groups": output_groups, } if has_billing_amounts: result["billing_overview"] = { field: fixed_decimal(value, precision) for field, value in billing_overview.items() } result["billing_overview_formula"] = ( "totalCost = runningTotal - discounts - credits" ) return result def write_json(path: Path, value: Any) -> None: path.parent.mkdir(parents=True, exist_ok=True) path.write_text(json.dumps(value, indent=2, ensure_ascii=False) + "\n", encoding="utf-8") def load_grouping_decision( path: Path, input_tenant_ids: set[str] ) -> tuple[bool, dict[str, Any]]: decision = json.loads(path.read_text(encoding="utf-8")) if not isinstance(decision, dict): raise ValueError("Grouping decision must be a JSON object.") if decision.get("format") != "manual-invoice-user-unit-cost-grouping-decision/v1": raise ValueError("Unsupported grouping decision format.") mode = decision.get("grouping_mode") if mode not in {"standard", "ignore_labels"}: raise ValueError("Grouping decision mode must be standard or ignore_labels.") expected_ignore_labels = mode == "ignore_labels" if decision.get("ignore_labels") is not expected_ignore_labels: raise ValueError( "Grouping decision ignore_labels does not match grouping_mode." ) tenant_id = str(decision.get("tenant_id", "")) if input_tenant_ids and input_tenant_ids != {tenant_id}: raise ValueError( "Grouping decision tenant does not match the input tenant IDs." ) return expected_ignore_labels, decision def main() -> int: args = parse_args() rows = list(load_jsonl(args.input)) grouping_decision: dict[str, Any] if args.grouping_decision: input_tenant_ids = { str(row["tenant_id"]) for _, row in rows if "tenant_id" in row } ignore_labels, grouping_decision = load_grouping_decision( args.grouping_decision, input_tenant_ids ) else: ignore_labels = args.ignore_labels grouping_decision = { "grouping_mode": "ignore_labels" if ignore_labels else "standard", "reason": "explicit_flag" if ignore_labels else "default_standard", } result = calculate( rows, amount_fields=args.amount_fields, precision=args.precision, ignore_invoice_id=args.ignore_invoice_id, ignore_labels=ignore_labels, include_billing_day=args.include_billing_day, rollup_fields=args.rollup_fields, ) result["grouping_decision"] = grouping_decision result["input_path"] = str(args.input) result["input_sha256"] = hashlib.sha256(args.input.read_bytes()).hexdigest() write_json(args.output, result) print( json.dumps( { key: value for key, value in result.items() if key not in {"groups", "rollups"} } ) ) return 0 if __name__ == "__main__": raise SystemExit(main())