# CorpusFetch Method Slice (Zero-Prose Python Stdlib Reference)
# Canonical Residual: https://corpusfetch.com/c/what-structural-signals-indicate-token-accounting-drift-in-opentelemetry-genai-agent-spans-when-tool-call-latency-overlaps-model-inference.md
# Remnant Warrant:    https://corpusfetch.com/c/what-structural-signals-indicate-token-accounting-drift-in-opentelemetry-genai-agent-spans-when-tool-call-latency-overlaps-model-inference.warrant.json
# Markdown SHA-256:   ef0ec8f02f0e34e7030ddd02841d90f6a228dabec033416484a45cf1adee9b27
# Code AST SHA-256:   7f66cd0077debb0efc15f17ac45f07b830e5b0c439634da5263d90737bae64a8
# Verified Classes:   TokenAttributionError, ChatSpan
# Defensive Guards:   1 fail-closed invariants (zero external deps)

from dataclasses import dataclass
from typing import Iterable


class TokenAttributionError(RuntimeError):
    """Raised when agent-span token usage cannot be attributed unambiguously."""


@dataclass(frozen=True)
class ChatSpan:
    span_id: str
    parent_id: str
    operation: str          # "chat" | "execute_tool" | "invoke_agent"
    start_ns: int
    end_ns: int
    input_tokens: int
    output_tokens: int
    response_id: str | None  # None => provider omitted usage object
    model: str | None


def classify_usage(spans: Iterable[ChatSpan]) -> dict[str, str]:
    """Return span_id -> 'measured' | 'inherited' | 'aggregate'.

    Raises TokenAttributionError when a chat span nested under a non-model
    tool span lacks a response_id, making attribution impossible.
    """
    spans = list(spans)
    by_id = {s.span_id: s for s in spans}
    labels: dict[str, str] = {}

    for s in spans:
        if s.operation == "invoke_agent" and (s.input_tokens or s.output_tokens):
            labels[s.span_id] = "aggregate"
            continue
        if s.operation != "chat":
            continue

        parent = by_id.get(s.parent_id)
        if parent is not None and parent.operation == "execute_tool":
            if s.response_id is None:
                # Nested chat under a tool, no provider usage object:
                # the usage block was inherited, not measured.
                sibling = next(
                    (o for o in spans
                     if o.operation == "chat"
                     and o.span_id != s.span_id
                     and o.input_tokens == s.input_tokens
                     and o.output_tokens == s.output_tokens),
                    None,
                )
                if sibling is None:
                    raise TokenAttributionError(
                        f"chat span {s.span_id} nested under tool {parent.span_id} "
                        f"has no response_id and no matching sibling; cannot "
                        f"decide measured vs inherited"
                    )
                labels[s.span_id] = "inherited"
            else:
                labels[s.span_id] = "measured"
        else:
            labels[s.span_id] = "measured" if s.response_id else "inherited"

    return labels


def total_attributed_tokens(spans: Iterable[ChatSpan]) -> int:
    spans = list(spans)
    labels = classify_usage(spans)
    aggregate = [s for s in spans if labels.get(s.span_id) == "aggregate"]
    if aggregate:
        # Parent invoke_agent usage is authoritative; ignore children.
        return sum(s.input_tokens + s.output_tokens for s in aggregate)
    return sum(
        s.input_tokens + s.output_tokens
        for s in spans
        if labels.get(s.span_id) == "measured"
    )