import { bigint, index, integer, jsonb, numeric, pgTable, text, timestamp, uuid } from "drizzle-orm/pg-core";
import { researchRuns } from "./runs.js";

export const providerUsage = pgTable(
  "provider_usage",
  {
    id: uuid("id").primaryKey().defaultRandom(),
    runId: uuid("run_id")
      .notNull()
      .references(() => researchRuns.id, { onDelete: "cascade" }),
    attemptId: uuid("attempt_id"),
    stage: text("stage"),
    provider: text("provider").notNull().default("anthropic"),
    model: text("model").notNull(),
    inputTokens: bigint("input_tokens", { mode: "number" }).notNull().default(0),
    outputTokens: bigint("output_tokens", { mode: "number" }).notNull().default(0),
    cacheReadTokens: bigint("cache_read_tokens", { mode: "number" }).notNull().default(0),
    cacheCreateTokens: bigint("cache_create_tokens", { mode: "number" }).notNull().default(0),
    costUsd: numeric("cost_usd", { precision: 10, scale: 4 }).notNull().default("0"),
    turns: integer("turns").notNull().default(0),
    createdAt: timestamp("created_at", { withTimezone: true }).notNull().defaultNow(),
  },
  (t) => [index("provider_usage_run_idx").on(t.runId)],
);

export const resourceUsage = pgTable("resource_usage", {
  id: uuid("id").primaryKey().defaultRandom(),
  runId: uuid("run_id")
    .notNull()
    .references(() => researchRuns.id, { onDelete: "cascade" }),
  attemptId: uuid("attempt_id"),
  workerId: uuid("worker_id"),
  cpuSeconds: numeric("cpu_seconds", { precision: 12, scale: 2 }),
  wallSeconds: numeric("wall_seconds", { precision: 12, scale: 2 }),
  diskBytes: bigint("disk_bytes", { mode: "number" }),
  peakRssBytes: bigint("peak_rss_bytes", { mode: "number" }),
  createdAt: timestamp("created_at", { withTimezone: true }).notNull().defaultNow(),
});

export const costRecords = pgTable("cost_records", {
  id: uuid("id").primaryKey().defaultRandom(),
  runId: uuid("run_id").references(() => researchRuns.id, { onDelete: "cascade" }),
  kind: text("kind").notNull(),
  amountUsd: numeric("amount_usd", { precision: 10, scale: 4 }).notNull(),
  detail: jsonb("detail").$type<Record<string, unknown>>().notNull().default({}),
  createdAt: timestamp("created_at", { withTimezone: true }).notNull().defaultNow(),
});