# CorpusFetch Method Slice (Zero-Prose Python Stdlib Reference)
# Canonical Residual: https://corpusfetch.com/c/how-should-an-instance-attribute-cost-and-latency-to-nested-agent-spans-when-genai-semantic-conventions-make-parent-and-child-span-accounting-non-addi.md
# Remnant Warrant:    https://corpusfetch.com/c/how-should-an-instance-attribute-cost-and-latency-to-nested-agent-spans-when-genai-semantic-conventions-make-parent-and-child-span-accounting-non-addi.warrant.json
# Markdown SHA-256:   d5413d6549a0865e1f66d6ee21962ce1b00c6a771c094dcf7e12dba6010024a7
# Code AST SHA-256:   037b6d4696833131ae95e14fd72cd952bd1fb277911f22609943559a0a35288a
# Verified Classes:   AttributionInvariantError, Span
# Defensive Guards:   5 fail-closed invariants (zero external deps)

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Iterable


class AttributionInvariantError(RuntimeError):
    """Raised when a trace cannot be reduced to a single billable total
    without violating the leaf-only usage invariant."""


@dataclass
class Span:
    span_id: str
    parent_id: str | None
    operation: str
    input_tokens: int = 0
    output_tokens: int = 0
    response_model: str | None = None
    request_model: str | None = None
    duration_ms: float = 0.0
    retry_of: str | None = None


PRICE_PER_1K = {"gpt-4o": 0.005, "gpt-4o-mini": 0.00015, "fallback-l": 0.0002}


def _price(span: Span) -> float:
    # response_model is authoritative; request_model is advisory only.
    model = span.response_model
    if model is None:
        if span.input_tokens or span.output_tokens:
            raise AttributionInvariantError(
                f"span {span.span_id} reports tokens but no gen_ai.response.model; "
                "cannot price without risking fallback rate drift"
            )
        return 0.0
    if model not in PRICE_PER_1K:
        raise AttributionInvariantError(f"unknown model rate for {model!r}")
    return (span.input_tokens + span.output_tokens) / 1000.0 * PRICE_PER_1K[model]


def reduce_trace(spans: Iterable[Span]) -> dict[str, float]:
    spans = list(spans)
    by_id = {s.span_id: s for s in spans}
    parents = {s.parent_id for s in spans if s.parent_id}

    # Detect upward usage propagation: a span that has children must not
    # carry usage copied from a descendant. We cannot see provenance, so we
    # enforce the structural rule: any span with children that also reports
    # tokens is flagged as parent-only overhead, and its children are leaves.
    leaf_cost = 0.0
    overhead_cost = 0.0
    seen_retry_targets: set[str] = set()

    for s in spans:
        if s.retry_of is not None:
            # Failed attempt: usage must be zero; do not bill it.
            if s.input_tokens or s.output_tokens:
                raise AttributionInvariantError(
                    f"retry span {s.span_id} carries usage for a retried attempt"
                )
            seen_retry_targets.add(s.retry_of)
            continue
        if s.span_id in seen_retry_targets:
            # A later successful attempt exists; bill once, at the winner.
            pass
        cost = _price(s)
        if s.span_id in parents:
            overhead_cost += cost
        else:
            leaf_cost += cost

    # Latency: parent wall-clock is authoritative; child durations only
    # attribute, never sum to the total.
    root_durations = [s.duration_ms for s in spans if s.parent_id is None]
    if len(root_durations) != 1:
        raise AttributionInvariantError(
            f"expected exactly one root span, found {len(root_durations)}"
        )
    end_to_end_ms = root_durations[0]
    child_sum_ms = sum(s.duration_ms for s in spans if s.parent_id is not None)
    if child_sum_ms > end_to_end_ms * 1.001:
        raise AttributionInvariantError(
            "child durations exceed parent wall-clock; retries or clock skew "
            "present — latency is non-additive and must use critical path"
        )

    return {
        "direct_model_cost": round(leaf_cost, 6),
        "orchestration_overhead": round(overhead_cost, 6),
        "total_cost": round(leaf_cost + overhead_cost, 6),
        "end_to_end_ms": end_to_end_ms,
        "unattributed_ms": round(end_to_end_ms - child_sum_ms, 3),
    }
