import { getSql } from "../connection";
import type { ProviderName } from "../../agent/models";

/**
 * Usage read from the message ledger rather than the session rollup.
 *
 * Every agent turn stores its own per-model usage, so the messages table is the
 * only place a model's tokens can still be told apart; the session totals
 * collapse them.
 *
 * `costUsd` is whatever the backend reported and is never recomputed from a
 * rate table — a derived figure presented as a bill would be fiction dressed as
 * a measurement. `estimatedCostUsd` is that derived figure, kept in its own
 * column and never summed into cost, because codex reporting nothing at all is
 * how a total Claude outage came to read as weekly spend falling to $0.00.
 */

export interface UsageRow {
  /** Local calendar day, `YYYY-MM-DD`. */
  day: string;
  room: string;
  model: string;
  provider: string;
  inputTokens: number;
  outputTokens: number;
  cacheReadTokens: number;
  cacheWriteTokens: number;
  /** null when nothing in this bucket carried a cost. */
  costUsd: number | null;
  /** List-rate projection for turns nobody billed. Never added to `costUsd`. */
  estimatedCostUsd: number | null;
  turns: number;
  unpricedTurns: number;
}

export interface UsageQuery {
  since: Date;
  room?: string;
  timezone?: string;
}

/** The Claude SDK names the deployment (firstParty/bedrock/vertex) where the
 *  chain names the backend. Both appear in stored history. */
const DEPLOYMENTS: Record<string, ProviderName> = {
  firstparty: "claude",
  bedrock: "claude",
  vertex: "claude",
};

export function normalizeProvider(raw: string | null | undefined): string {
  const value = (raw ?? "").trim();
  if (!value) return "unknown";
  return DEPLOYMENTS[value.toLowerCase()] ?? value;
}

const num = (v: unknown): number => {
  const n = typeof v === "string" ? Number(v) : v;
  return typeof n === "number" && Number.isFinite(n) ? n : 0;
};

export async function queryUsage({ since, room, timezone = "UTC" }: UsageQuery): Promise<UsageRow[]> {
  const sql = getSql();
  const rows = await sql<Record<string, unknown>[]>`
    SELECT
      to_char(m.created_at AT TIME ZONE ${timezone}, 'YYYY-MM-DD')          AS day,
      m.room                                                                AS room,
      COALESCE(NULLIF(u.value->>'canonicalModel', ''), u.key)               AS model,
      u.value->>'provider'                                                  AS provider,
      SUM(COALESCE((u.value->>'inputTokens')::bigint, 0))                   AS input_tokens,
      SUM(COALESCE((u.value->>'outputTokens')::bigint, 0))                  AS output_tokens,
      SUM(COALESCE((u.value->>'cacheReadInputTokens')::bigint, 0))          AS cache_read_tokens,
      SUM(COALESCE((u.value->>'cacheCreationInputTokens')::bigint, 0))      AS cache_write_tokens,
      SUM((u.value->>'costUSD')::numeric)                                   AS cost_usd,
      SUM((u.value->>'estimatedCostUSD')::numeric)                          AS estimated_cost_usd,
      COUNT(*)                                                              AS turns,
      COUNT(*) FILTER (WHERE u.value->>'costUSD' IS NULL)                   AS unpriced_turns
    FROM messages m
    CROSS JOIN LATERAL jsonb_each(m.metadata->'model_usage') AS u
    WHERE m.is_from_agent
      AND jsonb_typeof(m.metadata->'model_usage') = 'object'
      AND m.created_at >= ${since}
      ${room ? sql`AND m.room = ${room}` : sql``}
    GROUP BY 1, 2, 3, 4
  `;

  return merge(
    rows.map((r) => ({
      day: String(r.day),
      room: String(r.room),
      model: String(r.model),
      provider: normalizeProvider(r.provider as string | null),
      inputTokens: num(r.input_tokens),
      outputTokens: num(r.output_tokens),
      cacheReadTokens: num(r.cache_read_tokens),
      cacheWriteTokens: num(r.cache_write_tokens),
      costUsd: r.cost_usd === null ? null : num(r.cost_usd),
      estimatedCostUsd: r.estimated_cost_usd === null ? null : num(r.estimated_cost_usd),
      turns: num(r.turns),
      unpricedTurns: num(r.unpriced_turns),
    })),
  );
}

/** Normalizing the provider can collapse two SQL groups into one. */
function merge(rows: UsageRow[]): UsageRow[] {
  const byKey = new Map<string, UsageRow>();
  for (const row of rows) {
    const key = `${row.day} ${row.room} ${row.model} ${row.provider}`;
    const seen = byKey.get(key);
    if (!seen) {
      byKey.set(key, { ...row });
      continue;
    }
    seen.inputTokens += row.inputTokens;
    seen.outputTokens += row.outputTokens;
    seen.cacheReadTokens += row.cacheReadTokens;
    seen.cacheWriteTokens += row.cacheWriteTokens;
    seen.turns += row.turns;
    seen.unpricedTurns += row.unpricedTurns;
    seen.costUsd = seen.costUsd === null && row.costUsd === null ? null : (seen.costUsd ?? 0) + (row.costUsd ?? 0);
    seen.estimatedCostUsd =
      seen.estimatedCostUsd === null && row.estimatedCostUsd === null
        ? null
        : (seen.estimatedCostUsd ?? 0) + (row.estimatedCostUsd ?? 0);
  }
  return [...byKey.values()].sort((a, b) => a.day.localeCompare(b.day) || a.model.localeCompare(b.model));
}