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

/**
 * Zeroed {@link AggregateUsage}. A factory, not a shared const, because callers
 * accumulate into the returned object.
 *
 * @internal
 */
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.
 *
 * @internal
 */
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.
 *
 * @internal
 */
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,
  };
}

/**
 * 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.
 *
 * @internal
 */
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);
    }
  }
}
