# CorpusFetch Method Slice (Zero-Prose Python Stdlib Reference)
# Canonical Residual: https://corpusfetch.com/c/how-does-an-instance-detect-token-accounting-drift-when-opentelemetry-genai-agent-spans-split-a-single-model-invocation-across-tool-call-boundaries.md
# Remnant Warrant:    https://corpusfetch.com/c/how-does-an-instance-detect-token-accounting-drift-when-opentelemetry-genai-agent-spans-split-a-single-model-invocation-across-tool-call-boundaries.warrant.json
# Markdown SHA-256:   42cadf43714b7cd131e65234d753f84b0071602d914bfd7e13cba6f60ba82479
# Code AST SHA-256:   5434f3cb7e154c75e414c20d8745444cf6623bd3d8f34c6a42c93d6b0d48c316
# Verified Classes:   TokenAccountingDrift, Span
# Defensive Guards:   5 fail-closed invariants (zero external deps)

from __future__ import annotations

from dataclasses import dataclass
from typing import Iterable


class TokenAccountingDrift(RuntimeError):
    """Raised when span-level token sums cannot be reconciled to a billed total."""
    pass


@dataclass(frozen=True)
class Span:
    name: str
    operation: str            # gen_ai.operation.name
    input_tokens: int | None  # None == attribute absent, distinct from 0
    output_tokens: int | None
    parent_id: str | None
    tool_call_id: str | None = None


INVOCATION_OPS = {"chat", "text_completion", "invoke_agent"}
TOOL_OP = "execute_tool"


def classify_tree(spans: Iterable[Span]) -> str:
    spans = list(spans)
    parents = [s for s in spans if s.operation in INVOCATION_OPS and s.parent_id is None]
    children = [s for s in spans if s.operation in INVOCATION_OPS and s.parent_id is not None]
    parent_has_usage = any(
        p.input_tokens is not None or p.output_tokens is not None for p in parents
    )
    child_has_usage = any(
        c.input_tokens is not None or c.output_tokens is not None for c in children
    )
    if parent_has_usage and child_has_usage:
        return "duplicated_suspected"
    if parent_has_usage:
        return "flat_aggregate"
    if child_has_usage:
        return "nested_per_turn"
    return "unaccounted"


def authoritative_input_tokens(spans: Iterable[Span]) -> int:
    spans = list(spans)
    shape = classify_tree(spans)
    if shape == "flat_aggregate":
        return sum(
            s.input_tokens or 0
            for s in spans
            if s.operation in INVOCATION_OPS and s.parent_id is None
        )
    if shape == "nested_per_turn":
        return sum(
            s.input_tokens or 0
            for s in spans
            if s.operation in INVOCATION_OPS and s.parent_id is not None
        )
    if shape == "duplicated_suspected":
        # Cannot decide without external reference; refuse to guess.
        raise TokenAccountingDrift(
            "parent and child invocation spans both carry usage; "
            "shape=duplicated_suspected, reconcile against billed total before summing"
        )
    raise TokenAccountingDrift("no usage attributes on any invocation span; shape=unaccounted")


def reconcile(
    spans: Iterable[Span],
    billed_input_tokens: int,
    turn_count: int,
    tolerance: float = 0.02,
) -> float:
    """Return the observed/billed ratio; raise if it matches duplication signature."""
    spans = list(spans)
    try:
        observed = authoritative_input_tokens(spans)
    except TokenAccountingDrift:
        # Duplicated shape: fall back to parent-only sum for the ratio test.
        observed = sum(
            s.input_tokens or 0
            for s in spans
            if s.operation in INVOCATION_OPS and s.parent_id is None
        )
    if billed_input_tokens <= 0:
        raise TokenAccountingDrift("billed_input_tokens must be positive for ratio test")
    ratio = observed / billed_input_tokens
    if abs(ratio - 1.0) <= tolerance:
        return ratio
    if turn_count > 1 and abs(ratio - turn_count) / turn_count <= tolerance:
        raise TokenAccountingDrift(
            f"ratio={ratio:.3f} matches turn_count={turn_count}; "
            "producer double-counts usage across parent and child spans"
        )
    raise TokenAccountingDrift(
        f"ratio={ratio:.3f} unexplained; neither ~1.0 nor ~turn_count={turn_count}"
    )


def dedupe_tool_spans(spans: Iterable[Span], consumed_ids: set[str]) -> list[Span]:
    seen: dict[str, Span] = {}
    for s in spans:
        if s.operation != TOOL_OP or s.tool_call_id is None:
            continue
        if s.tool_call_id in consumed_ids:
            seen[s.tool_call_id] = s
        elif s.tool_call_id not in seen:
            seen[s.tool_call_id] = s  # keep last-emitted if never consumed
    return list(seen.values())