// Metered usage, persisted to the database.
//
// `appendUsage` in enterprise.ts writes JSONL to the data directory, which is right for a
// self-hosted box and useless for a hosted one: Cloud Run's filesystem is ephemeral and scales to
// zero, so that file disappears regularly. Nothing can be billed or rate-limited on top of it.
// This is the same row, in Postgres (or sqlite when self-hosting), which survives a restart.
//
// What's stored is the raw fact — tokens, model, provider — not a cost. Prices change; tokens
// don't. A cost column would freeze whatever price table happened to be loaded on the day, and
// would have to be recomputed anyway the first time a provider repriced. Cost is derived at read
// time from the model id.
import type { Pool } from "pg";
import type Database from "better-sqlite3";
import { authDatabase, usingPostgres } from "./db.js";
import { priceOf } from "../client/models-dev.js";

export interface UsageEvent {
  ts: number;
  user: string;
  model: string;
  provider: string;
  promptTokens: number;
  completionTokens: number;
  /** Wall time of the whole response, and time to first token. Both were already measured at the
   *  call site and handed to the JSONL writer; only this table dropped them, which meant latency
   *  could not be analysed on a hosted box at all (that file is ephemeral on Cloud Run). */
  ms?: number;
  ttftMs?: number;
  /** IANA zone the CLIENT reported (x-ada-tz), e.g. "Asia/Kolkata". Coarse by construction and
   *  chosen deliberately over an IP lookup: it answers both questions we actually have — roughly
   *  where usage comes from, and what the local hour was — without storing an address. */
  tz?: string;
  /** Two-letter country, only when a proxy in front of us already resolved one. We never look it
   *  up ourselves and never store the IP it came from. */
  country?: string;
}

const pg = () => authDatabase() as Pool;
const lite = () => authDatabase() as Database.Database;

let ready: Promise<void> | null = null;
function ensure(): Promise<void> {
  ready ??= (async () => {
    // (user_id, ts) is the only access pattern that matters: "what has this account spent this
    // period". Without the index that becomes a full scan on the busiest table in the system.
    const ddl = usingPostgres
      ? `create table if not exists usage_events (
           id bigserial primary key,
           ts bigint not null,
           user_id text not null,
           model text not null,
           provider text not null,
           prompt_tokens integer not null default 0,
           completion_tokens integer not null default 0,
           ms integer,
           ttft_ms integer,
           tz text,
           country text
         )`
      : `create table if not exists usage_events (
           id integer primary key autoincrement,
           ts integer not null,
           user_id text not null,
           model text not null,
           provider text not null,
           prompt_tokens integer not null default 0,
           completion_tokens integer not null default 0,
           ms integer,
           ttft_ms integer,
           tz text,
           country text
         )`;
    const idx = "create index if not exists usage_events_user_ts on usage_events (user_id, ts)";
    // Columns added after the table shipped. `create table if not exists` is a no-op against an
    // existing table, so without this every deployment that predates them keeps the old shape and
    // every insert fails on the unknown column — the metering write is best-effort, so that failure
    // would be silent and permanent. All four are nullable: rows written before this stay valid,
    // and analytics treats null as "not captured" rather than zero.
    const added: Array<[string, string]> = [
      ["ms", "integer"],
      ["ttft_ms", "integer"],
      ["tz", "text"],
      ["country", "text"],
    ];
    if (usingPostgres) {
      await pg().query(ddl);
      await pg().query(idx);
      for (const [col, type] of added) await pg().query(`alter table usage_events add column if not exists ${col} ${type}`);
    } else {
      lite().exec(ddl);
      lite().exec(idx);
      // sqlite has no ADD COLUMN IF NOT EXISTS — ask the table what it already has.
      const have = new Set((lite().prepare("pragma table_info(usage_events)").all() as Array<{ name: string }>).map((r) => r.name));
      for (const [col, type] of added) if (!have.has(col)) lite().exec(`alter table usage_events add column ${col} ${type}`);
    }
  })();
  return ready;
}

/** Record one metered response. Best-effort by design — the same contract as the file writer: a
 *  metering failure must never fail the user's request. Failures are logged, not thrown, because
 *  this is called from inside a response-stream teardown where there is nobody left to catch. */
export async function recordUsage(e: UsageEvent): Promise<void> {
  try {
    await ensure();
    if (usingPostgres) {
      await pg().query(
        "insert into usage_events (ts, user_id, model, provider, prompt_tokens, completion_tokens, ms, ttft_ms, tz, country) values ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)",
        [e.ts, e.user, e.model, e.provider, e.promptTokens, e.completionTokens, e.ms ?? null, e.ttftMs ?? null, e.tz ?? null, e.country ?? null],
      );
    } else {
      lite()
        .prepare(
          "insert into usage_events (ts, user_id, model, provider, prompt_tokens, completion_tokens, ms, ttft_ms, tz, country) values (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
        )
        .run(e.ts, e.user, e.model, e.provider, e.promptTokens, e.completionTokens, e.ms ?? null, e.ttftMs ?? null, e.tz ?? null, e.country ?? null);
    }
  } catch (err) {
    console.error("[ada] usage write failed:", err instanceof Error ? err.message : err);
  }
}

export interface UsageTotal {
  promptTokens: number;
  completionTokens: number;
  requests: number;
}

/** What one account has used since a timestamp — the number a quota is checked against. */
export async function usageSince(user: string, sinceMs: number): Promise<UsageTotal> {
  await ensure();
  const sql =
    "select coalesce(sum(prompt_tokens),0) as p, coalesce(sum(completion_tokens),0) as c, count(*) as n from usage_events where user_id = $1 and ts >= $2";
  const row = usingPostgres
    ? ((await pg().query(sql, [user, sinceMs])).rows[0] as { p: string; c: string; n: string })
    : (lite().prepare(sql.replace(/\$\d/g, "?")).get(user, sinceMs) as { p: number; c: number; n: number });
  return { promptTokens: Number(row?.p ?? 0), completionTokens: Number(row?.c ?? 0), requests: Number(row?.n ?? 0) };
}

/** $ per 1M [input, output] tokens. `:free` costs nothing. Anything models.dev doesn't price gets a
 *  deliberately PESSIMISTIC guess: a cap that under-counts an unknown model is a cap that doesn't
 *  cap, and an unpriced id is exactly where a surprise bill comes from. */
const UNKNOWN_PRICE: [number, number] = [3, 15];
export function priceUsd(model: string): [number, number] {
  if (/:free$/i.test(model)) return [0, 0];
  return priceOf(model) ?? UNKNOWN_PRICE;
}

export interface Spend extends UsageTotal {
  usd: number;
}

/** What one account has COST us since a timestamp — the number the spend cap is checked against.
 *
 *  Derived at read time from stored tokens × today's price, not stored per row: prices move, and a
 *  frozen cost column would have to be recomputed the first time a provider repriced anyway. Free
 *  models fall out for free — they price at 0, so they add 0 without a special case. */
export async function costSince(user: string, sinceMs: number): Promise<Spend> {
  const rows = await usageByModel(user, sinceMs);
  let usd = 0;
  const total: Spend = { usd: 0, promptTokens: 0, completionTokens: 0, requests: 0 };
  for (const r of rows) {
    const [inPrice, outPrice] = priceUsd(r.model);
    usd += (r.promptTokens * inPrice + r.completionTokens * outPrice) / 1_000_000;
    total.promptTokens += r.promptTokens;
    total.completionTokens += r.completionTokens;
    total.requests += r.requests;
  }
  total.usd = usd;
  return total;
}

/** Per-model breakdown for an account over a window — for a usage page, and for costing a period
 *  once a price table exists (cost needs the model, which is why it's stored per row). */
export async function usageByModel(user: string, sinceMs: number): Promise<Array<UsageTotal & { model: string }>> {
  await ensure();
  const sql =
    "select model, coalesce(sum(prompt_tokens),0) as p, coalesce(sum(completion_tokens),0) as c, count(*) as n from usage_events where user_id = $1 and ts >= $2 group by model order by n desc";
  const rows = usingPostgres
    ? ((await pg().query(sql, [user, sinceMs])).rows as Array<{ model: string; p: string; c: string; n: string }>)
    : (lite().prepare(sql.replace(/\$\d/g, "?")).all(user, sinceMs) as Array<{ model: string; p: number; c: number; n: number }>);
  return rows.map((r) => ({ model: r.model, promptTokens: Number(r.p), completionTokens: Number(r.c), requests: Number(r.n) }));
}