import type { RunManifest } from "@autoscholar/schemas";
import {
  bigint,
  bigserial,
  boolean,
  index,
  integer,
  jsonb,
  numeric,
  pgTable,
  text,
  timestamp,
  uniqueIndex,
  uuid,
} from "drizzle-orm/pg-core";
import { adminUsers } from "./auth.js";
import { projectInstructionVersions, researchProjects } from "./projects.js";
import { computeWorkers } from "./workers.js";

export const researchRuns = pgTable(
  "research_runs",
  {
    id: uuid("id").primaryKey().defaultRandom(),
    projectId: uuid("project_id")
      .notNull()
      .references(() => researchProjects.id),
    instructionVersionId: uuid("instruction_version_id")
      .notNull()
      .references(() => projectInstructionVersions.id),
    idempotencyKey: text("idempotency_key").notNull(),
    title: text("title").notNull(),
    status: text("status").notNull().default("draft"),
    researchMode: text("research_mode").notNull(),
    secondaryModes: text("secondary_modes").array().notNull().default([]),
    discipline: text("discipline").notNull(),
    computeProfile: text("compute_profile").notNull().default("standard"),
    provider: text("provider").notNull().default("anthropic"),
    model: text("model").notNull(),
    reviewerModel: text("reviewer_model").notNull(),
    maxBudgetUsd: numeric("max_budget_usd", { precision: 10, scale: 2 }).notNull(),
    maxRuntimeMinutes: integer("max_runtime_minutes").notNull(),
    maxAttempts: integer("max_attempts").notNull().default(2),
    autoPublish: boolean("auto_publish").notNull().default(false),
    qualityPolicyVersion: text("quality_policy_version").notNull(),
    manifest: jsonb("manifest").$type<RunManifest>().notNull(),
    manifestSha256: text("manifest_sha256").notNull(),
    currentAttemptNo: integer("current_attempt_no").notNull().default(0),
    currentStage: text("current_stage"),
    fencingToken: bigint("fencing_token", { mode: "number" }).notNull().default(0),
    lastEventSeq: bigint("last_event_seq", { mode: "number" }).notNull().default(0),
    cancelRequested: boolean("cancel_requested").notNull().default(false),
    failureReason: text("failure_reason"),
    totalCostUsd: numeric("total_cost_usd", { precision: 10, scale: 4 }).notNull().default("0"),
    paperId: uuid("paper_id"),
    queuedAt: timestamp("queued_at", { withTimezone: true }),
    startedAt: timestamp("started_at", { withTimezone: true }),
    finishedAt: timestamp("finished_at", { withTimezone: true }),
    createdBy: uuid("created_by").references(() => adminUsers.id),
    createdAt: timestamp("created_at", { withTimezone: true }).notNull().defaultNow(),
    updatedAt: timestamp("updated_at", { withTimezone: true }).notNull().defaultNow(),
  },
  (t) => [
    uniqueIndex("runs_idempotency_uq").on(t.idempotencyKey),
    index("runs_status_idx").on(t.status, t.queuedAt),
    index("runs_project_idx").on(t.projectId),
  ],
);

export const runAttempts = pgTable(
  "run_attempts",
  {
    id: uuid("id").primaryKey().defaultRandom(),
    runId: uuid("run_id")
      .notNull()
      .references(() => researchRuns.id, { onDelete: "cascade" }),
    attemptNo: integer("attempt_no").notNull(),
    workerId: uuid("worker_id").references(() => computeWorkers.id),
    status: text("status").notNull().default("leased"), // leased | running | completed | failed | cancelled | expired | budget_exceeded | timeout
    fencingToken: bigint("fencing_token", { mode: "number" }).notNull(),
    resumeFromStage: text("resume_from_stage"),
    startedAt: timestamp("started_at", { withTimezone: true }).notNull().defaultNow(),
    endedAt: timestamp("ended_at", { withTimezone: true }),
    error: text("error"),
    costUsd: numeric("cost_usd", { precision: 10, scale: 4 }).notNull().default("0"),
    summary: jsonb("summary").$type<Record<string, unknown>>(),
  },
  (t) => [uniqueIndex("attempts_run_no_uq").on(t.runId, t.attemptNo), index("attempts_run_idx").on(t.runId)],
);

export const stageExecutions = pgTable(
  "stage_executions",
  {
    id: uuid("id").primaryKey().defaultRandom(),
    runId: uuid("run_id")
      .notNull()
      .references(() => researchRuns.id, { onDelete: "cascade" }),
    attemptId: uuid("attempt_id")
      .notNull()
      .references(() => runAttempts.id, { onDelete: "cascade" }),
    stage: text("stage").notNull(),
    status: text("status").notNull().default("pending"),
    startedAt: timestamp("started_at", { withTimezone: true }).notNull().defaultNow(),
    endedAt: timestamp("ended_at", { withTimezone: true }),
    costUsd: numeric("cost_usd", { precision: 10, scale: 4 }).notNull().default("0"),
    turns: integer("turns").notNull().default(0),
    usage: jsonb("usage").$type<Record<string, unknown>>(),
    error: text("error"),
  },
  (t) => [
    uniqueIndex("stage_exec_attempt_stage_uq").on(t.attemptId, t.stage),
    index("stage_exec_run_idx").on(t.runId),
  ],
);

export const stageEvents = pgTable(
  "stage_events",
  {
    id: bigserial("id", { mode: "number" }).primaryKey(),
    runId: uuid("run_id")
      .notNull()
      .references(() => researchRuns.id, { onDelete: "cascade" }),
    attemptId: uuid("attempt_id"),
    workerId: uuid("worker_id"),
    seq: bigint("seq", { mode: "number" }).notNull(),
    stage: text("stage"),
    level: text("level").notNull().default("info"),
    kind: text("kind").notNull(),
    message: text("message").notNull(),
    data: jsonb("data").$type<Record<string, unknown>>(),
    at: timestamp("at", { withTimezone: true }).notNull(),
    createdAt: timestamp("created_at", { withTimezone: true }).notNull().defaultNow(),
  },
  (t) => [uniqueIndex("stage_events_run_seq_uq").on(t.runId, t.seq)],
);

export const jobLeases = pgTable("job_leases", {
  runId: uuid("run_id")
    .primaryKey()
    .references(() => researchRuns.id, { onDelete: "cascade" }),
  attemptId: uuid("attempt_id").notNull(),
  workerId: uuid("worker_id")
    .notNull()
    .references(() => computeWorkers.id),
  fencingToken: bigint("fencing_token", { mode: "number" }).notNull(),
  acquiredAt: timestamp("acquired_at", { withTimezone: true }).notNull().defaultNow(),
  expiresAt: timestamp("expires_at", { withTimezone: true }).notNull(),
  lastHeartbeatAt: timestamp("last_heartbeat_at", { withTimezone: true }).notNull().defaultNow(),
});

export const checkpoints = pgTable(
  "checkpoints",
  {
    id: uuid("id").primaryKey().defaultRandom(),
    runId: uuid("run_id")
      .notNull()
      .references(() => researchRuns.id, { onDelete: "cascade" }),
    attemptId: uuid("attempt_id").notNull(),
    stage: text("stage").notNull(),
    data: jsonb("data").$type<Record<string, unknown>>().notNull().default({}),
    createdAt: timestamp("created_at", { withTimezone: true }).notNull().defaultNow(),
  },
  (t) => [index("checkpoints_run_idx").on(t.runId, t.createdAt)],
);