import type { AggregateUsage, ChatMessage } from '@genesislcap/foundation-ai';

/**
 * Zeroed {@link AggregateUsage}. A factory, not a shared const, because callers
 * accumulate into the returned object.
 *
 * @beta
 */
export function emptyUsage(): AggregateUsage {
  return {
    costUsd: 0,
    uncachedInputTokens: 0,
    cacheReadTokens: 0,
    cacheWriteTokens: 0,
    outputTokens: 0,
  };
}

/**
 * Total tokens processed — the four buckets added up. Safe to sum because the
 * buckets are disjoint (see {@link AggregateUsage}); the equivalent sum over
 * per-message fields would double-count the prompt.
 *
 * Bear in mind what this number is: because every turn resends the conversation,
 * a turn's prompt is counted again in the next turn's. It is billed throughput,
 * NOT the count of distinct tokens in the conversation, so it grows
 * super-linearly with turn count and is expected to dwarf the context size.
 *
 * @beta
 */
export function totalTokens(usage: AggregateUsage): number {
  return (
    usage.uncachedInputTokens + usage.cacheReadTokens + usage.cacheWriteTokens + usage.outputTokens
  );
}

/**
 * `a + b`, per bucket, as a new object.
 *
 * Used to lay a row's banked (pre-transcript) spend over what the live transcript
 * proves. Pure rather than in-place because both operands are live state a caller
 * must not mutate.
 *
 * @beta
 */
export function addUsage(a: AggregateUsage, b: AggregateUsage): AggregateUsage {
  return {
    costUsd: a.costUsd + b.costUsd,
    uncachedInputTokens: a.uncachedInputTokens + b.uncachedInputTokens,
    cacheReadTokens: a.cacheReadTokens + b.cacheReadTokens,
    cacheWriteTokens: a.cacheWriteTokens + b.cacheWriteTokens,
    outputTokens: a.outputTokens + b.outputTokens,
  };
}

/**
 * The four buckets plus USD for ONE message — the per-message fields projected into
 * the disjoint, addable shape, or `undefined` when the message reports no usage at
 * all (a user turn, a display-only reasoning/narration split, a provider that
 * reports none).
 *
 * Shallow by design: it does NOT recurse into `toolCalls[].subAgentTrace` and does
 * NOT count `compaction.rolledUpUsage`, so it answers "what did this one request
 * cost" rather than "what did this message and everything under it cost". Use
 * {@link sumUsage} for the latter — passing a whole transcript through this one
 * message at a time would silently drop both.
 *
 * `costUsd` folds in `externalCostUsd` (non-LLM spend a widget reported), matching
 * `sumUsage`. It is `0` when the provider priced nothing, which is NOT a claim that
 * the call was free — see {@link AggregateUsage}; `UsageRow.costUsd` is the
 * shape that keeps unpriced distinguishable, and is the right tool for a ledger.
 *
 * @beta
 */
export function messageUsage(m: ChatMessage): AggregateUsage | undefined {
  if (
    m.cost == null &&
    m.externalCostUsd == null &&
    m.inputTokens == null &&
    m.outputTokens == null &&
    m.cacheReadTokens == null &&
    m.cacheWriteTokens == null
  ) {
    return undefined;
  }
  const cacheReadTokens = m.cacheReadTokens ?? 0;
  const cacheWriteTokens = m.cacheWriteTokens ?? 0;
  return {
    costUsd: (m.cost ?? 0) + (m.externalCostUsd ?? 0),
    // The uncached REMAINDER, clamped — the same arithmetic `accumulate` below applies,
    // for the same reasons. The two are deliberately separate implementations (that walk
    // mutates one accumulator rather than allocating per message) and must agree;
    // `sum-usage.test.ts` pins them together.
    uncachedInputTokens: Math.max(0, (m.inputTokens ?? 0) - cacheReadTokens - cacheWriteTokens),
    cacheReadTokens,
    cacheWriteTokens,
    outputTokens: m.outputTokens ?? 0,
  };
}

/**
 * Sum cost and per-bucket token usage across a message list.
 *
 * One walk produces everything a usage display needs, replacing separate cost and
 * token passes over the same (potentially deep) transcript. Recurses into each
 * tool call's `subAgentTrace`, so LLM work a sub-agent did — possibly on another
 * provider at other rates — lands in the total; the driver embeds a whole
 * sub-agent conversation in the parent's tool result, so this is the only place
 * that spend is visible.
 *
 * Splitting the prompt into buckets is the point: `cost` already has the cache
 * discount applied, but a single token total cannot show that a session was cheap
 * because most of its prompt was a cache hit. The three input buckets are derived
 * from the per-message fields, where `inputTokens` is the whole prompt and the
 * cache fields break it down — so uncached input is the remainder, never an
 * addition.
 *
 * Also counts `compaction.rolledUpUsage`: compaction destroys the messages it
 * summarises, and without their banked usage every total re-derived from the
 * transcript would drop by whatever those turns had already spent.
 *
 * Returns a zeroed total when no message carries usage — providers that report
 * none (e.g. Chrome built-in) are skipped silently rather than guessed at.
 *
 * @remarks
 * This sums the WHOLE array, which is the right answer only when the array is the
 * whole of the work being billed. A driver whose conversation **continues across
 * rounds** opens its history with turns an earlier round already paid for, so
 * summing it again bills them twice — and the session total then grows
 * quadratically in the number of rounds while every individual figure in it stays
 * correct. Nothing about the result looks wrong, which is why it is called out
 * here rather than left to the caller to notice.
 *
 * Do NOT solve this with a "sum from this message onwards" marker. Compaction
 * *replaces* the messages it summarises — and compaction is one of this function's
 * own call sites — so the marker message can be absent from history by the time
 * the sum runs, leaving the caller to either fail or silently sum everything,
 * which is the double-count the marker was meant to prevent.
 *
 * Instead hold a **banked baseline**: spend that predates the current transcript,
 * carried as a value, with `usage = banked + live` composed via {@link addUsage}.
 * `resolveBankedUsage` in `utils/cost-session-history.ts` is the worked form, and
 * the shape of its signature is the load-bearing part. It separates three cases:
 *
 * 1. no prior record — a genuine first round, so nothing is banked;
 * 2. prior record, and the live transcript already accounts for it (a restored
 *    session) — carry the prior record's own banked figure forward;
 * 3. prior record, and the transcript starts empty — bank the prior record's
 *    whole total, because nothing in the transcript can re-derive it.
 *
 * Case 1 and case 3 take the same arithmetic and mean opposite things, so the
 * prior record must be a **required** input with no zero-defaulting path: a
 * caller who carries a transcript and forgets its baseline reports the entire
 * prior conversation as this round's spend. Distinguishing 2 from 3 is a separate
 * decision, and getting it wrong is destructive in the other direction — it
 * rewrites a total down to the current page-load and drops earlier spend.
 *
 * @beta
 */
export function sumUsage(messages: readonly ChatMessage[]): AggregateUsage {
  const total = emptyUsage();
  accumulate(messages, total);
  return total;
}

/** Add `from` into `into`, in place. */
function addInto(into: AggregateUsage, from: AggregateUsage): void {
  into.costUsd += from.costUsd;
  into.uncachedInputTokens += from.uncachedInputTokens;
  into.cacheReadTokens += from.cacheReadTokens;
  into.cacheWriteTokens += from.cacheWriteTokens;
  into.outputTokens += from.outputTokens;
}

/**
 * Walk `messages`, accumulating into `into`. Mutates a single accumulator rather
 * than merging returned objects — this runs on every transcript change, so the
 * per-message allocation is worth avoiding.
 */
function accumulate(messages: readonly ChatMessage[], into: AggregateUsage): void {
  for (const m of messages) {
    // Provider/LLM cost plus any non-LLM cost a widget reported for its own
    // external service calls. Both are real spend on this session.
    if (m.cost != null) into.costUsd += m.cost;
    if (m.externalCostUsd != null) into.costUsd += m.externalCostUsd;

    const cacheRead = m.cacheReadTokens ?? 0;
    const cacheWrite = m.cacheWriteTokens ?? 0;
    into.cacheReadTokens += cacheRead;
    into.cacheWriteTokens += cacheWrite;
    // The uncached remainder of the prompt. Clamped at zero so malformed or
    // hand-edited history whose buckets exceed the prompt total can't drive the
    // session total negative — a wrong-but-plausible number is worse than a
    // conservative one here, since this feeds a cost display.
    into.uncachedInputTokens += Math.max(0, (m.inputTokens ?? 0) - cacheRead - cacheWrite);
    into.outputTokens += m.outputTokens ?? 0;

    // Spend banked by a compaction on this summary's behalf. NOTE: deliberately
    // not `compaction.tokensBefore` — that is a one-turn context-window reading,
    // not cumulative spend, and counting it here would inflate the total.
    const rolled = m.compaction?.rolledUpUsage;
    if (rolled) addInto(into, rolled);

    for (const tc of m.toolCalls ?? []) {
      if (tc.subAgentTrace) accumulate(tc.subAgentTrace, into);
    }
  }
}
