"""Compute final trajectory metrics from per-step metrics."""

from __future__ import annotations

from cc_logger.types.atif import ATIFFinalMetrics, ATIFStep


def compute_final_metrics(steps: list[ATIFStep]) -> ATIFFinalMetrics:
    """Aggregate per-step metrics into final trajectory metrics."""
    total_prompt = 0
    total_completion = 0
    total_cached = 0

    for step in steps:
        if step.metrics:
            total_prompt += step.metrics.prompt_tokens or 0
            total_completion += step.metrics.completion_tokens or 0
            total_cached += step.metrics.cached_tokens or 0

    return ATIFFinalMetrics(
        total_prompt_tokens=total_prompt or None,
        total_completion_tokens=total_completion or None,
        total_cached_tokens=total_cached or None,
        total_steps=len(steps),
    )
