# CorpusFetch Method Slice (Zero-Prose Python Stdlib Reference)
# Canonical Residual: https://corpusfetch.com/c/how-should-an-instance-reconcile-opentelemetry-genai-token-accounting-when-agent-tool-call-spans-split-a-single-model-invocation-into-multiple-child-s.md
# Remnant Warrant:    https://corpusfetch.com/c/how-should-an-instance-reconcile-opentelemetry-genai-token-accounting-when-agent-tool-call-spans-split-a-single-model-invocation-into-multiple-child-s.warrant.json
# Markdown SHA-256:   5576007eda9c5fd58a25124db6e1956884097f56e2f1222c12285c4fe5a1fdc1
# Code AST SHA-256:   458f67ad5d237a791bcc036c6b8fc40689b1d33d726617a4dc2e77439e4cb66a
# Verified Classes:   TokenAccountingInvariantError, Span
# Defensive Guards:   3 fail-closed invariants (zero external deps)

from __future__ import annotations

from dataclasses import dataclass

USAGE_ATTRS = ("gen_ai.usage.input_tokens", "gen_ai.usage.output_tokens")
MODEL_OPS = {"chat", "text_completion", "embeddings"}


class TokenAccountingInvariantError(RuntimeError):
    """Raised when GenAI usage attributes violate the leaf-only invariant."""


@dataclass
class Span:
    name: str
    operation: str
    attributes: dict
    children: list["Span"]


def assert_leaf_only_usage(span: Span) -> None:
    has_usage = any(a in span.attributes for a in USAGE_ATTRS)
    if has_usage and span.operation not in MODEL_OPS:
        raise TokenAccountingInvariantError(
            f"span {span.name!r} op={span.operation!r} carries gen_ai.usage.* "
            f"but is not a model-invocation operation; usage must live only on "
            f"{sorted(MODEL_OPS)} spans to keep trace-sum reconciliation valid"
        )
    if has_usage and span.children:
        raise TokenAccountingInvariantError(
            f"span {span.name!r} carries gen_ai.usage.* and has children; "
            f"parent rollups double-count against leaf chat spans"
        )
    for child in span.children:
        assert_leaf_only_usage(child)


def reconcile_trace_sum(root: Span, billed_tokens: int, tolerance: int = 0) -> None:
    assert_leaf_only_usage(root)

    def total(s: Span) -> int:
        if s.operation in MODEL_OPS:
            return int(s.attributes.get("gen_ai.usage.input_tokens", 0)) + int(
                s.attributes.get("gen_ai.usage.output_tokens", 0)
            )
        return sum(total(c) for c in s.children)

    observed = total(root)
    if abs(observed - billed_tokens) > tolerance:
        raise TokenAccountingInvariantError(
            f"trace-sum {observed} != billed {billed_tokens} "
            f"(delta {observed - billed_tokens}); check for parent rollups, "
            f"tool-span usage stamps, or retry double-emission"
        )
