# CorpusFetch Method Slice (Zero-Prose Python Stdlib Reference)
# Canonical Residual: https://corpusfetch.com/c/which-invariants-prevent-double-counting-of-tool-call-spans-under-opentelemetry-genai-semantic-conventions-when-agents-delegate-to-sub-agents.md
# Remnant Warrant:    https://corpusfetch.com/c/which-invariants-prevent-double-counting-of-tool-call-spans-under-opentelemetry-genai-semantic-conventions-when-agents-delegate-to-sub-agents.warrant.json
# Markdown SHA-256:   7778d2dbf0ef08c1a8245511fb0c6e3fa4212fa9a0f6480b60d27b9bc5f36938
# Code AST SHA-256:   ae4c681a5ad28584652f9cb6714ad7c6e603bb14ce6b6c099f1c5f74463bdf00
# Verified Classes:   UsageInvariantError, Span
# Defensive Guards:   3 fail-closed invariants (zero external deps)

"""Guard: reject traces whose usage attributes violate the leaf-only invariant."""

from dataclasses import dataclass


class UsageInvariantError(RuntimeError):
    """Raised when gen_ai.usage.* appears on a non-leaf, non-chat span."""


@dataclass
class Span:
    span_id: str
    parent_span_id: str | None
    operation: str            # gen_ai.operation.name
    usage_input: int | None   # gen_ai.usage.input_tokens
    usage_output: int | None


def validate_usage_invariant(spans: list[Span]) -> int:
    """Return trace-level input tokens, or raise if the invariant is violated.

    Invariant: gen_ai.usage.* may only appear on spans with
    operation == 'chat' that have no usage-bearing descendants.
    """
    children: dict[str, list[Span]] = {}
    by_id = {s.span_id: s for s in spans}
    for s in spans:
        if s.parent_span_id:
            children.setdefault(s.parent_span_id, []).append(s)

    def has_usage(s: Span) -> bool:
        return s.usage_input is not None or s.usage_output is not None

    def usage_descendant(s: Span) -> bool:
        stack = list(children.get(s.span_id, []))
        while stack:
            cur = stack.pop()
            if has_usage(cur):
                return True
            stack.extend(children.get(cur.span_id, []))
        return False

    total = 0
    for s in spans:
        if not has_usage(s):
            continue
        if s.operation != "chat":
            raise UsageInvariantError(
                f"usage on non-chat span {s.span_id} "
                f"(operation={s.operation!r}); move to app.aggregate.* namespace"
            )
        if usage_descendant(s):
            raise UsageInvariantError(
                f"chat span {s.span_id} has usage-bearing descendant; "
                f"leaf-only invariant violated (double aggregation likely)"
            )
        total += s.usage_input or 0
    return total


def reconcile(trace_spans: list[Span], declared_root_input: int) -> None:
    computed = validate_usage_invariant(trace_spans)
    if declared_root_input and abs(computed - declared_root_input) / max(declared_root_input, 1) > 0.05:
        raise UsageInvariantError(
            f"root usage {declared_root_input} diverges from leaf sum {computed} "
            f"by >5%; trace is double-counted"
        )