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

/**
 * One billable unit of spend, projected out of a transcript.
 *
 * @remarks
 * Deliberately **not** a `ChatMessage`, and not an extension of one. A row is a
 * derived record for a ledger, and keeping the types distinct makes the one
 * dangerous mistake — feeding rows back into `sumUsage`, which does its own
 * recursion and would count everything twice — a compile error rather than a
 * plausible wrong number.
 *
 * @beta
 */
export interface UsageRow {
  /**
   * Where the spend came from.
   *
   * - `'request'` — one model call, with a message behind it.
   * - `'compaction'` — spend **banked** by a compaction. The turns that incurred it
   *   were summarised and deleted, so there is no single call to attribute it to;
   *   the amount survives on the summary's `compaction.rolledUpUsage`. Emitted as
   *   its own row because otherwise this money is simply missing from a per-call
   *   view while remaining present in the total, and the two disagree silently.
   */
  source: 'request' | 'compaction';
  /**
   * Model that produced the call. `undefined` on a `'compaction'` row: the banked
   * spend covers however many turns were summarised, which may have run on
   * different models, so attributing it to one would invent a fact. (The summary's
   * `compaction.model` is the model that did the *summarising* — a different thing,
   * and not what this spend was.)
   */
  model?: string;
  /**
   * Vendor that served the call. Taken from the provider the driver stamped on the message,
   * which is authoritative and covers vendors the model allowlists do not — a server-proxied
   * `gpt-5` turn is `'openai'` here and unknown to `vendorOfModel`. Falls back to resolving
   * {@link UsageRow.model} through `vendorOfModel` only for messages persisted before that
   * field existed, and is left unset when neither knows rather than guessed at.
   */
  provider?: AIProviderType;
  /**
   * USD for this row, as reported by the provider and stamped by the transport.
   *
   * `undefined` means no cost was reported — **not** zero. Nothing here estimates
   * one from the token counts: an estimate is plausible, unremarkable and quietly
   * wrong, whereas a total that is visibly too low gets investigated. A row with
   * tokens but no `costUsd` is worth alarming on; that is the shape that check needs.
   */
  costUsd?: number;
  /** Non-LLM spend a host reported for this turn (an external service it called). */
  externalCostUsd?: number;
  /** Prompt tokens billed at the full uncached rate — the remainder after the cache buckets. */
  uncachedInputTokens: number;
  /** Prompt tokens served from the provider's cache. */
  cacheReadTokens: number;
  /** Prompt tokens written to the provider's cache. */
  cacheWriteTokens: number;
  /** Generated tokens, including reasoning tokens the provider bills as output. */
  outputTokens: number;
  /** Agent that produced the call, when known. */
  agentName?: string;
  /** Delegation depth: 0 for the top-level conversation, 1 for a sub-agent, 2 for its sub-agent, … */
  subAgentDepth: number;
  /** Id of the tool call that spawned this row's conversation — correlates a row back to its delegation. */
  subAgentOf?: string;
}

/** Whether a message carries anything a ledger would record. */
function hasUsage(m: ChatMessage): boolean {
  return (
    m.cost != null ||
    m.externalCostUsd != null ||
    m.inputTokens != null ||
    m.outputTokens != null ||
    m.cacheReadTokens != null ||
    m.cacheWriteTokens != null
  );
}

/**
 * Project a transcript into one row per billable unit of spend.
 *
 * @remarks
 * Reads the messages and returns a table; it stores nothing, mutates nothing, and
 * adds nothing to history. Use it where an **aggregate is not enough** — writing
 * usage rows for per-project attribution, a cost dashboard, or auditing which turns
 * came back unpriced. For a single total, `sumUsage` is the answer and this is the
 * wrong tool.
 *
 * Walks the same shape `sumUsage` does, and is guaranteed to agree with it: summing
 * `costUsd` and `externalCostUsd` across these rows equals `sumUsage(...).costUsd`
 * over the same input, and the four token buckets reconcile likewise. That property
 * is the point of the function — a per-call view that quietly disagrees with the
 * total is worse than no per-call view, because both look right in isolation.
 *
 * Covers the two places spend hides:
 *
 * - **Sub-agent conversations.** A delegating tool call carries the child's entire
 *   conversation on `toolCall.subAgentTrace`, so a walk of top-level messages alone
 *   misses everything a delegating agent spent. A consumer that reimplemented this
 *   and omitted the branch reported $0.106 on a run that cost $0.234.
 * - **Compactions.** A compaction deletes the messages it summarises, banking their
 *   spend on the summary. Those turns produce a single `source: 'compaction'` row.
 *
 * Messages carrying no usage at all (a user turn, a narration) produce no row.
 *
 * **Do not pass the result to `sumUsage`** — the types prevent it, and the reason is
 * that both recurse, so a pre-flattened list would be counted twice. Both functions
 * take raw history.
 *
 * @beta
 */
export function usageRows(messages: readonly ChatMessage[]): UsageRow[] {
  return collectRows(messages, 0, undefined);
}

/** Map an already-bucketed {@link AggregateUsage} onto a row's token fields. */
function bucketsOf(
  usage: AggregateUsage,
): Pick<UsageRow, 'uncachedInputTokens' | 'cacheReadTokens' | 'cacheWriteTokens' | 'outputTokens'> {
  return {
    uncachedInputTokens: usage.uncachedInputTokens,
    cacheReadTokens: usage.cacheReadTokens,
    cacheWriteTokens: usage.cacheWriteTokens,
    outputTokens: usage.outputTokens,
  };
}

function collectRows(
  messages: readonly ChatMessage[],
  depth: number,
  subAgentOf: string | undefined,
): UsageRow[] {
  const rows: UsageRow[] = [];
  for (const m of messages) {
    if (hasUsage(m)) {
      const cacheReadTokens = m.cacheReadTokens ?? 0;
      const cacheWriteTokens = m.cacheWriteTokens ?? 0;
      rows.push({
        source: 'request',
        model: m.model,
        // The driver stamps `ChatMessage.provider` from the resolved provider's own
        // status, so it is AUTHORITATIVE and covers vendors the model allowlists do
        // not — a server-proxied `gpt-5` turn is `provider: 'openai'` on the message
        // and unknown to `vendorOfModel`. Prefer it; fall back to resolving the model
        // id only for messages persisted before the field existed. Still left
        // undefined when neither knows, rather than guessed at.
        provider: m.provider ?? (m.model ? vendorOfModel(m.model) : undefined),
        ...(m.cost != null && { costUsd: m.cost }),
        ...(m.externalCostUsd != null && { externalCostUsd: m.externalCostUsd }),
        // Uncached input is the REMAINDER, matching `sumUsage`: `inputTokens` is the
        // whole prompt and the cache fields break it down. Clamped for the same
        // reason — hand-edited history must not drive a total negative.
        uncachedInputTokens: Math.max(0, (m.inputTokens ?? 0) - cacheReadTokens - cacheWriteTokens),
        cacheReadTokens,
        cacheWriteTokens,
        outputTokens: m.outputTokens ?? 0,
        ...(m.agentName != null && { agentName: m.agentName }),
        subAgentDepth: depth,
        ...(subAgentOf != null && { subAgentOf }),
      });
    }

    // Spend the compaction banked on this summary's behalf. Its buckets are already
    // disjoint (it is an `AggregateUsage`), so they map across without the
    // subtraction a per-message record needs.
    const rolled = m.compaction?.rolledUpUsage;
    if (rolled) {
      rows.push({
        source: 'compaction',
        ...(rolled.costUsd != null && { costUsd: rolled.costUsd }),
        ...bucketsOf(rolled),
        ...(m.agentName != null && { agentName: m.agentName }),
        subAgentDepth: depth,
        ...(subAgentOf != null && { subAgentOf }),
      });
    }

    for (const tc of m.toolCalls ?? []) {
      if (tc.subAgentTrace) rows.push(...collectRows(tc.subAgentTrace, depth + 1, tc.id));
    }
  }
  return rows;
}
