# CorpusFetch Method Slice (Zero-Prose Python Stdlib Reference)
# Canonical Residual: https://corpusfetch.com/c/where-must-an-instance-of-a-genai-agent-span-end-before-the-parent-llm-span-s-token-accounting-closes-the-tool-call-re-entry-boundary.md
# Remnant Warrant:    https://corpusfetch.com/c/where-must-an-instance-of-a-genai-agent-span-end-before-the-parent-llm-span-s-token-accounting-closes-the-tool-call-re-entry-boundary.warrant.json
# Markdown SHA-256:   418cddc563bb398280b6a3cbd0d365b45aa961604a3da979e4b955a03d8291dd
# Code AST SHA-256:   649c8b852e79dc1949d1aa2bd47dff96542eabdbc1762e8d8acfa1123dc84127
# Verified Classes:   GenAIAccountingBoundaryError, SpanRecord, AgentLoop
# Defensive Guards:   6 fail-closed invariants (zero external deps)

from dataclasses import dataclass, field


class GenAIAccountingBoundaryError(RuntimeError):
    """Raised when token usage would be attributed across a span-lifetime
    boundary that violates the per-invocation accounting invariant."""


@dataclass
class SpanRecord:
    span_id: str
    operation: str  # "invoke_agent" | "chat" | "execute_tool"
    parent_id: str | None
    input_tokens: int = 0
    output_tokens: int = 0
    closed: bool = False


@dataclass
class AgentLoop:
    agent_span_id: str
    llm_spans: list[SpanRecord] = field(default_factory=list)
    tool_spans: list[SpanRecord] = field(default_factory=list)

    def open_llm_span(self, span_id: str) -> SpanRecord:
        open_llm = [s for s in self.llm_spans if not s.closed]
        if open_llm:
            raise GenAIAccountingBoundaryError(
                f"cannot open LLM span {span_id!r}: prior LLM span "
                f"{open_llm[0].span_id!r} is still open; usage would be "
                f"double-counted or lost on tool re-entry"
            )
        span = SpanRecord(span_id, "chat", self.agent_span_id)
        self.llm_spans.append(span)
        return span

    def close_with_usage(self, span_id: str, inp: int, out: int) -> None:
        match = [s for s in self.llm_spans if s.span_id == span_id]
        if not match:
            raise GenAIAccountingBoundaryError(
                f"usage reported for unknown LLM span {span_id!r}"
            )
        span = match[0]
        if span.closed:
            raise GenAIAccountingBoundaryError(
                f"span {span_id!r} already closed; usage would overwrite"
            )
        if span.operation != "chat":
            raise GenAIAccountingBoundaryError(
                f"gen_ai.usage.* may only be set on 'chat' spans, got "
                f"{span.operation!r}"
            )
        span.input_tokens = inp
        span.output_tokens = out
        span.closed = True

    def open_tool_span(self, span_id: str) -> SpanRecord:
        open_llm = [s for s in self.llm_spans if not s.closed]
        if open_llm:
            raise GenAIAccountingBoundaryError(
                f"tool span {span_id!r} opened while LLM span "
                f"{open_llm[0].span_id!r} is open; tool must follow a "
                f"closed LLM invocation"
            )
        span = SpanRecord(span_id, "execute_tool", self.agent_span_id)
        self.tool_spans.append(span)
        return span

    def agent_rollup(self) -> tuple[int, int]:
        if any(not s.closed for s in self.llm_spans):
            raise GenAIAccountingBoundaryError(
                "cannot emit agent rollup: LLM spans still open, "
                "loop has not terminated"
            )
        return (
            sum(s.input_tokens for s in self.llm_spans),
            sum(s.output_tokens for s in self.llm_spans),
        )