agent-sessions 0.2.1

use crate::{
    AccountingPolicy, Endpoint, LineErrorKind, TokenSemantics, Usage, UsageAdjustment, UsageBasis,
};
use crate::{
    accounting::counts,
    parser::{State, string},
};
use serde_json::Value;

pub(super) fn has_unaccounted_last(p: &Value) -> bool {
    let Some(info) = p.get("info") else {
        return false;
    };
    if info.get("total_token_usage").is_some_and(|v| !v.is_null()) {
        return false;
    }
    let Some(last) = info.get("last_token_usage") else {
        return false;
    };
    [
        "input_tokens",
        "output_tokens",
        "cached_input_tokens",
        "cache_read_input_tokens",
        "cache_write_input_tokens",
        "reasoning_output_tokens",
        "total_tokens",
    ]
    .into_iter()
    .filter_map(|field| last.get(field).and_then(Value::as_u64))
    .chain(
        last.pointer("/cache_creation/ephemeral_1h_input_tokens")
            .and_then(Value::as_u64),
    )
    .any(|n| n > 0)
}

pub(super) fn token_count(
    p: &Value,
    state: &mut State,
    policy: AccountingPolicy,
) -> Result<Option<Usage>, LineErrorKind> {
    let statistics = policy == AccountingPolicy::UsageStatistics;
    let info = match p.get("info") {
        Some(Value::Null) => return Ok(None),
        None if statistics || p.get("rate_limits").is_some() => return Ok(None),
        Some(i) if i.is_object() => i,
        _ => return Err(LineErrorKind::MissingField("info")),
    };
    let total = match info.get("total_token_usage").filter(|v| !v.is_null()) {
        Some(total) => total,
        None if statistics => return Ok(None),
        None => return Err(LineErrorKind::MissingField("total_token_usage")),
    };
    let (total, adjustments) = counts(total, false, policy)?;
    let last = info
        .get("last_token_usage")
        .filter(|v| !v.is_null())
        .map(|v| counts(v, false, policy))
        .transpose()?;
    from_counts(total, last, adjustments, state, policy)
}

pub(crate) fn from_counts(
    total: crate::TokenCounts,
    last: Option<(crate::TokenCounts, Vec<UsageAdjustment>)>,
    mut adjustments: Vec<UsageAdjustment>,
    state: &mut State,
    policy: AccountingPolicy,
) -> Result<Option<Usage>, LineErrorKind> {
    let statistics = policy == AccountingPolicy::UsageStatistics;
    let duplicate = state.previous.is_some_and(|prev| {
        if statistics {
            total.statistics_eq(prev)
        } else {
            total == prev
        }
    });
    if (!state.broken_usage || statistics) && duplicate {
        return Ok(None);
    }
    let previous = state.previous.replace(total);
    let broken = std::mem::replace(&mut state.broken_usage, false);
    let (counts, basis) = if let Some((last, changes)) = last {
        adjustments.extend(changes);
        (last, UsageBasis::LastSample)
    } else {
        if broken {
            if !statistics {
                return Err(LineErrorKind::LostUsageBaseline);
            }
            adjustments.push(UsageAdjustment::UncertainCumulativeBaseline);
        }
        if let Some(prev) = previous {
            if total.regressed_from(&prev) {
                if !statistics {
                    return Err(LineErrorKind::CounterRegression);
                }
                adjustments.push(UsageAdjustment::CumulativeRegression);
            }
            if statistics {
                if total
                    .values()
                    .into_iter()
                    .zip(prev.values())
                    .any(|(a, b)| a.is_some() && b.is_none())
                {
                    adjustments.push(UsageAdjustment::MissingCumulativeBaseline);
                }
                (total.statistics_delta(prev), UsageBasis::CumulativeDelta)
            } else {
                (total.delta(prev), UsageBasis::CumulativeDelta)
            }
        } else {
            (total, UsageBasis::InitialCumulative)
        }
    };
    if counts.is_missing() && !statistics {
        return Err(LineErrorKind::LostUsageBaseline);
    }
    Ok(Some(Usage {
        adjustments,
        dedup_key: None,
        model: None,
        counts,
        cumulative: Some(total),
        semantics: TokenSemantics::CodexInclusive,
        basis,
        stop_reason: None,
        endpoint: Endpoint::Unknown,
        inference_geo: None,
    }))
}

pub(super) fn response(p: &Value, state: &State) -> Result<Usage, LineErrorKind> {
    let (counts, adjustments) = counts(
        p.get("usage").ok_or(LineErrorKind::MissingField("usage"))?,
        false,
        AccountingPolicy::Strict,
    )?;
    Ok(Usage {
        adjustments,
        dedup_key: string(p, "response_id").map(|id| format!("codex:{id}")),
        model: string(p, "model")
            .filter(|s| !s.trim().is_empty())
            .or_else(|| state.model.clone()),
        counts,
        cumulative: None,
        semantics: TokenSemantics::CodexInclusive,
        basis: UsageBasis::Response,
        stop_reason: None,
        endpoint: Endpoint::Unknown,
        inference_geo: None,
    })
}

pub(crate) fn apply_model(usage: &mut Usage, state: &mut State, observed: Option<String>) {
    if let Some(model) = observed {
        if state.model.as_deref() != Some(model.as_str()) {
            state.model = Some(model.clone());
        }
        usage.model = Some(model);
    } else {
        usage.model = state.model.clone();
    }
}