image-prompt-studio/lib/db.ts

import fs from "node:fs";
import path from "node:path";
import { randomUUID } from "node:crypto";
import bcrypt from "bcryptjs";
import Database from "better-sqlite3";
import { appConfig } from "./config";
import {
  buildStoredImageDownloadUrl,
  buildStoredImageUrl,
} from "@/lib/image-storage";
import type {
  ImageHistoryItem,
  GenerationJobStatus,
  SessionRecord,
  SettingsSnapshot,
  StoredGeneratedImageInput,
  StoredImageAsset,
  UsageEventMeta,
  UserRecord,
  UserRole,
  UserUsageSummary,
} from "@/lib/types";

type SqliteDatabase = InstanceType<typeof Database>;

interface UserRow {
  id: string;
  username: string;
  password_hash: string;
  role: UserRole;
  monthly_quota: number | null;
  is_active: number;
  created_at: string;
  updated_at: string;
}

interface UserSummaryRow extends UserRow {
  used_this_period?: number | null;
}

interface SettingRow {
  key: string;
  value: string;
}

interface UsageRow {
  used_units?: number | null;
}

interface TableInfoRow {
  name: string;
}

interface StoredImageAssetRow {
  id: string;
  user_id: string;
  file_name: string;
  mime_type: string;
  file_size: number;
  storage_key: string;
}

interface ImageHistoryRow extends StoredImageAssetRow {
  job_id: string;
  username: string;
  prompt: string;
  revised_prompt: string | null;
  model: string | null;
  size: string | null;
  quality: string | null;
  output_format: string | null;
  background: string | null;
  moderation: string | null;
  created_at: string;
}

declare global {
  var __imagePromptStudioDb: SqliteDatabase | undefined;
}

function initializeDatabase(): SqliteDatabase {
  const databaseDirectory = path.dirname(appConfig.databasePath);

  fs.mkdirSync(databaseDirectory, { recursive: true });

  const database = new Database(appConfig.databasePath);

  database.pragma("journal_mode = WAL");
  database.pragma("foreign_keys = ON");

  database.exec(`
    CREATE TABLE IF NOT EXISTS users (
      id TEXT PRIMARY KEY,
      username TEXT NOT NULL UNIQUE,
      password_hash TEXT NOT NULL,
      role TEXT NOT NULL CHECK(role IN ('admin', 'user')),
      monthly_quota INTEGER,
      is_active INTEGER NOT NULL DEFAULT 1,
      created_at TEXT NOT NULL,
      updated_at TEXT NOT NULL
    );

    CREATE TABLE IF NOT EXISTS sessions (
      id TEXT PRIMARY KEY,
      user_id TEXT NOT NULL,
      expires_at TEXT NOT NULL,
      created_at TEXT NOT NULL,
      FOREIGN KEY(user_id) REFERENCES users(id) ON DELETE CASCADE
    );

    CREATE TABLE IF NOT EXISTS settings (
      key TEXT PRIMARY KEY,
      value TEXT NOT NULL,
      updated_at TEXT NOT NULL
    );

    CREATE TABLE IF NOT EXISTS usage_events (
      id TEXT PRIMARY KEY,
      user_id TEXT NOT NULL,
      units INTEGER NOT NULL,
      period_key TEXT NOT NULL,
      created_at TEXT NOT NULL,
      meta_json TEXT,
      FOREIGN KEY(user_id) REFERENCES users(id) ON DELETE CASCADE
    );

    CREATE TABLE IF NOT EXISTS generation_jobs (
      id TEXT PRIMARY KEY,
      user_id TEXT NOT NULL,
      prompt TEXT NOT NULL,
      status TEXT NOT NULL CHECK(status IN ('success', 'error')),
      image_count INTEGER NOT NULL DEFAULT 0,
      duration_ms INTEGER,
      model TEXT,
      size TEXT,
      quality TEXT,
      output_format TEXT,
      background TEXT,
      moderation TEXT,
      output_compression INTEGER,
      error_message TEXT,
      created_at TEXT NOT NULL,
      FOREIGN KEY(user_id) REFERENCES users(id) ON DELETE CASCADE
    );

    CREATE TABLE IF NOT EXISTS generated_images (
      id TEXT PRIMARY KEY,
      job_id TEXT NOT NULL,
      user_id TEXT NOT NULL,
      file_name TEXT NOT NULL,
      storage_key TEXT NOT NULL UNIQUE,
      mime_type TEXT NOT NULL,
      file_size INTEGER NOT NULL,
      revised_prompt TEXT,
      created_at TEXT NOT NULL,
      FOREIGN KEY(job_id) REFERENCES generation_jobs(id) ON DELETE CASCADE,
      FOREIGN KEY(user_id) REFERENCES users(id) ON DELETE CASCADE
    );

    CREATE INDEX IF NOT EXISTS idx_sessions_user_id ON sessions(user_id);
    CREATE INDEX IF NOT EXISTS idx_sessions_expires_at ON sessions(expires_at);
    CREATE INDEX IF NOT EXISTS idx_usage_events_user_period ON usage_events(user_id, period_key);
    CREATE INDEX IF NOT EXISTS idx_generation_jobs_user_id ON generation_jobs(user_id);
    CREATE INDEX IF NOT EXISTS idx_generated_images_user_created_at ON generated_images(user_id, created_at DESC);
    CREATE INDEX IF NOT EXISTS idx_generated_images_job_id ON generated_images(job_id);
  `);

  ensureColumnExists(database, "generation_jobs", "model", "TEXT");
  ensureColumnExists(database, "generation_jobs", "size", "TEXT");
  ensureColumnExists(database, "generation_jobs", "quality", "TEXT");
  ensureColumnExists(database, "generation_jobs", "output_format", "TEXT");
  ensureColumnExists(database, "generation_jobs", "background", "TEXT");
  ensureColumnExists(database, "generation_jobs", "moderation", "TEXT");
  ensureColumnExists(
    database,
    "generation_jobs",
    "output_compression",
    "INTEGER",
  );

  upsertSettingRecord(
    database,
    "default_monthly_quota",
    String(appConfig.defaultMonthlyQuota),
  );
  ensureAdminUser(database);

  return database;
}

function ensureAdminUser(database: SqliteDatabase): void {
  const now = new Date().toISOString();
  const existingAdmin = database
    .prepare("SELECT id FROM users WHERE username = ?")
    .get(appConfig.adminUsername) as { id: string } | undefined;

  const passwordHash = bcrypt.hashSync(appConfig.adminPassword, 12);

  if (existingAdmin) {
    if (!appConfig.adminSyncOnBoot) {
      return;
    }

    database
      .prepare(
        `
        UPDATE users
        SET
          password_hash = ?,
          role = 'admin',
          monthly_quota = NULL,
          is_active = 1,
          updated_at = ?
        WHERE username = ?
      `,
      )
      .run(passwordHash, now, appConfig.adminUsername);

    return;
  }

  database
    .prepare(
      `
      INSERT OR IGNORE INTO users (
        id,
        username,
        password_hash,
        role,
        monthly_quota,
        is_active,
        created_at,
        updated_at
      ) VALUES (?, ?, ?, 'admin', NULL, 1, ?, ?)
    `,
    )
    .run(randomUUID(), appConfig.adminUsername, passwordHash, now, now);
}

function upsertSettingRecord(
  database: SqliteDatabase,
  key: string,
  value: string,
): void {
  database
    .prepare(
      `
      INSERT INTO settings (key, value, updated_at)
      VALUES (?, ?, ?)
      ON CONFLICT(key) DO UPDATE SET
        value = excluded.value,
        updated_at = excluded.updated_at
    `,
    )
    .run(key, value, new Date().toISOString());
}

function ensureColumnExists(
  database: SqliteDatabase,
  tableName: string,
  columnName: string,
  columnSql: string,
): void {
  const columns = database
    .prepare(`PRAGMA table_info(${tableName})`)
    .all() as TableInfoRow[];

  if (columns.some((column) => column.name === columnName)) {
    return;
  }

  database.exec(
    `ALTER TABLE ${tableName} ADD COLUMN ${columnName} ${columnSql}`,
  );
}

if (!globalThis.__imagePromptStudioDb) {
  globalThis.__imagePromptStudioDb = initializeDatabase();
}

export const db = globalThis.__imagePromptStudioDb as SqliteDatabase;

export function syncAdminUserFromConfig(): void {
  ensureAdminUser(db);
}

export function normalizeUsername(value: unknown): string {
  return String(value ?? "")
    .trim()
    .toLowerCase();
}

export function getUserById(id: string): UserRecord | null {
  const row = db.prepare("SELECT * FROM users WHERE id = ?").get(id) as
    | UserRow
    | undefined;

  return mapUserRecord(row);
}

export function getUserByUsername(username: string): UserRecord | null {
  const row = db
    .prepare("SELECT * FROM users WHERE username = ?")
    .get(normalizeUsername(username)) as UserRow | undefined;

  return mapUserRecord(row);
}

export function listUserSummaries(periodKey: string): UserUsageSummary[] {
  const rows = db
    .prepare(
      `
      SELECT
        u.*,
        COALESCE(SUM(CASE WHEN ue.period_key = ? THEN ue.units ELSE 0 END), 0) AS used_this_period
      FROM users u
      LEFT JOIN usage_events ue ON ue.user_id = u.id
      GROUP BY u.id
      ORDER BY CASE u.role WHEN 'admin' THEN 0 ELSE 1 END, u.created_at ASC
    `,
    )
    .all(periodKey) as UserSummaryRow[];

  return rows.map((row) => {
    const user = mapUserRecord(row);

    if (!user) {
      throw new Error("Failed to map user record.");
    }

    return {
      ...user,
      usedThisPeriod: Number(row.used_this_period ?? 0),
    };
  });
}

export function createUserRecord({
  username,
  passwordHash,
  role,
  monthlyQuota,
}: {
  username: string;
  passwordHash: string;
  role: UserRole;
  monthlyQuota: number | string | null | undefined;
}): UserRecord {
  const now = new Date().toISOString();
  const id = randomUUID();

  db.prepare(
    `
    INSERT INTO users (
      id,
      username,
      password_hash,
      role,
      monthly_quota,
      is_active,
      created_at,
      updated_at
    ) VALUES (?, ?, ?, ?, ?, 1, ?, ?)
  `,
  ).run(
    id,
    normalizeUsername(username),
    passwordHash,
    role,
    normalizeMonthlyQuota(monthlyQuota),
    now,
    now,
  );

  const user = getUserById(id);

  if (!user) {
    throw new Error("Failed to create user record.");
  }

  return user;
}

export function updateUserRecord(
  id: string,
  updates: {
    role?: UserRole;
    monthlyQuota?: number | string | null;
    isActive?: boolean;
    passwordHash?: string;
  },
): UserRecord | null {
  const currentUser = getUserById(id);

  if (!currentUser) {
    return null;
  }

  const nextUser = {
    ...currentUser,
    role: updates.role ?? currentUser.role,
    monthlyQuota:
      updates.monthlyQuota === undefined
        ? currentUser.monthlyQuota
        : normalizeMonthlyQuota(updates.monthlyQuota),
    isActive:
      updates.isActive === undefined
        ? currentUser.isActive
        : Boolean(updates.isActive),
    passwordHash: updates.passwordHash ?? currentUser.passwordHash,
  };

  db.prepare(
    `
    UPDATE users
    SET
      role = ?,
      monthly_quota = ?,
      is_active = ?,
      password_hash = ?,
      updated_at = ?
    WHERE id = ?
  `,
  ).run(
    nextUser.role,
    nextUser.monthlyQuota,
    nextUser.isActive ? 1 : 0,
    nextUser.passwordHash,
    new Date().toISOString(),
    id,
  );

  return getUserById(id);
}

export function createSessionRecord({
  sessionId,
  userId,
  expiresAt,
}: SessionRecord): void {
  db.prepare(
    `
    INSERT INTO sessions (id, user_id, expires_at, created_at)
    VALUES (?, ?, ?, ?)
  `,
  ).run(sessionId, userId, expiresAt, new Date().toISOString());
}

export function deleteSessionRecord(sessionId: string): void {
  db.prepare("DELETE FROM sessions WHERE id = ?").run(sessionId);
}

export function deleteExpiredSessions(): void {
  db.prepare("DELETE FROM sessions WHERE expires_at <= ?").run(
    new Date().toISOString(),
  );
}

export function getSessionUser(sessionId: string): UserRecord | null {
  deleteExpiredSessions();

  const row = db
    .prepare(
      `
      SELECT u.*
      FROM sessions s
      INNER JOIN users u ON u.id = s.user_id
      WHERE s.id = ? AND s.expires_at > ?
      LIMIT 1
    `,
    )
    .get(sessionId, new Date().toISOString()) as UserRow | undefined;

  return mapUserRecord(row);
}

export function getSettingsSnapshot(): SettingsSnapshot {
  const rows = db
    .prepare("SELECT key, value FROM settings")
    .all() as SettingRow[];
  const settings = Object.fromEntries(rows.map((row) => [row.key, row.value]));

  return {
    defaultMonthlyQuota: Number.parseInt(
      settings.default_monthly_quota ?? String(appConfig.defaultMonthlyQuota),
      10,
    ),
  };
}

export function setDefaultMonthlyQuota(
  value: number | string,
): SettingsSnapshot {
  upsertSettingRecord(
    db,
    "default_monthly_quota",
    String(normalizeQuotaInput(value)),
  );

  return getSettingsSnapshot();
}

export function getUsageForUser(userId: string, periodKey: string): number {
  const row = db
    .prepare(
      `
      SELECT COALESCE(SUM(units), 0) AS used_units
      FROM usage_events
      WHERE user_id = ? AND period_key = ?
    `,
    )
    .get(userId, periodKey) as UsageRow | undefined;

  return Number(row?.used_units ?? 0);
}

export function recordUsageEvent({
  userId,
  units,
  periodKey,
  meta,
}: {
  userId: string;
  units: number;
  periodKey: string;
  meta?: UsageEventMeta | null;
}): void {
  insertUsageEvent(db, {
    userId,
    units,
    periodKey,
    meta,
    createdAt: new Date().toISOString(),
  });
}

export function recordCompletedGeneration({
  jobId = randomUUID(),
  userId,
  prompt,
  imageCount,
  durationMs,
  model,
  size,
  quality,
  outputFormat,
  background,
  moderation,
  outputCompression = null,
  periodKey,
  usageMeta,
  images,
}: {
  jobId?: string;
  userId: string;
  prompt: string;
  imageCount: number;
  durationMs: number | null;
  model: string;
  size: string;
  quality: string;
  outputFormat: string;
  background: string;
  moderation: string;
  outputCompression?: number | null;
  periodKey: string;
  usageMeta?: UsageEventMeta | null;
  images: StoredGeneratedImageInput[];
}): string {
  const now = new Date().toISOString();
  const commitGeneration = db.transaction(() => {
    insertGenerationJob(db, {
      id: jobId,
      userId,
      prompt,
      status: "success",
      imageCount,
      durationMs,
      errorMessage: null,
      model,
      size,
      quality,
      outputFormat,
      background,
      moderation,
      outputCompression,
      createdAt: now,
    });
    insertGeneratedImages(db, {
      jobId,
      userId,
      images,
      createdAt: now,
    });
    insertUsageEvent(db, {
      userId,
      units: imageCount,
      periodKey,
      meta: usageMeta,
      createdAt: now,
    });
  });

  commitGeneration();

  return jobId;
}

export function listImageHistoryForViewer({
  viewerId,
  viewerRole,
}: {
  viewerId: string;
  viewerRole: UserRole;
}): ImageHistoryItem[] {
  const rows =
    viewerRole === "admin"
      ? (db
          .prepare(
            `
            SELECT
              gi.id,
              gi.job_id,
              gi.user_id,
              gi.file_name,
              gi.storage_key,
              gi.mime_type,
              gi.file_size,
              gi.revised_prompt,
              gi.created_at,
              u.username,
              gj.prompt,
              gj.model,
              gj.size,
              gj.quality,
              gj.output_format,
              gj.background,
              gj.moderation
            FROM generated_images gi
            INNER JOIN generation_jobs gj ON gj.id = gi.job_id
            INNER JOIN users u ON u.id = gi.user_id
            ORDER BY gi.created_at DESC
          `,
          )
          .all() as ImageHistoryRow[])
      : (db
          .prepare(
            `
            SELECT
              gi.id,
              gi.job_id,
              gi.user_id,
              gi.file_name,
              gi.storage_key,
              gi.mime_type,
              gi.file_size,
              gi.revised_prompt,
              gi.created_at,
              u.username,
              gj.prompt,
              gj.model,
              gj.size,
              gj.quality,
              gj.output_format,
              gj.background,
              gj.moderation
            FROM generated_images gi
            INNER JOIN generation_jobs gj ON gj.id = gi.job_id
            INNER JOIN users u ON u.id = gi.user_id
            WHERE gi.user_id = ?
            ORDER BY gi.created_at DESC
          `,
          )
          .all(viewerId) as ImageHistoryRow[]);

  return rows.map(mapImageHistoryItem);
}

export function getStoredImageAssetForViewer({
  viewerId,
  viewerRole,
  imageId,
}: {
  viewerId: string;
  viewerRole: UserRole;
  imageId: string;
}): StoredImageAsset | null {
  const row =
    viewerRole === "admin"
      ? (db
          .prepare(
            `
            SELECT
              id,
              user_id,
              file_name,
              mime_type,
              file_size,
              storage_key
            FROM generated_images
            WHERE id = ?
            LIMIT 1
          `,
          )
          .get(imageId) as StoredImageAssetRow | undefined)
      : (db
          .prepare(
            `
            SELECT
              id,
              user_id,
              file_name,
              mime_type,
              file_size,
              storage_key
            FROM generated_images
            WHERE id = ? AND user_id = ?
            LIMIT 1
          `,
          )
          .get(imageId, viewerId) as StoredImageAssetRow | undefined);

  if (!row) {
    return null;
  }

  return {
    id: row.id,
    userId: row.user_id,
    fileName: row.file_name,
    mimeType: row.mime_type,
    fileSize: Number(row.file_size),
    storageKey: row.storage_key,
  };
}

export function recordGenerationJob({
  userId,
  prompt,
  status,
  imageCount = 0,
  durationMs = null,
  errorMessage = null,
  model = null,
  size = null,
  quality = null,
  outputFormat = null,
  background = null,
  moderation = null,
  outputCompression = null,
}: {
  userId: string;
  prompt: string;
  status: GenerationJobStatus;
  imageCount?: number;
  durationMs?: number | null;
  errorMessage?: string | null;
  model?: string | null;
  size?: string | null;
  quality?: string | null;
  outputFormat?: string | null;
  background?: string | null;
  moderation?: string | null;
  outputCompression?: number | null;
}): string {
  const id = randomUUID();

  insertGenerationJob(db, {
    id,
    userId,
    prompt,
    status,
    imageCount,
    durationMs,
    errorMessage,
    model,
    size,
    quality,
    outputFormat,
    background,
    moderation,
    outputCompression,
    createdAt: new Date().toISOString(),
  });

  return id;
}

function insertUsageEvent(
  database: SqliteDatabase,
  {
    userId,
    units,
    periodKey,
    meta,
    createdAt,
  }: {
    userId: string;
    units: number;
    periodKey: string;
    meta?: UsageEventMeta | null;
    createdAt: string;
  },
): void {
  database.prepare(
    `
    INSERT INTO usage_events (id, user_id, units, period_key, created_at, meta_json)
    VALUES (?, ?, ?, ?, ?, ?)
  `,
  ).run(
    randomUUID(),
    userId,
    units,
    periodKey,
    createdAt,
    meta ? JSON.stringify(meta) : null,
  );
}

function insertGenerationJob(
  database: SqliteDatabase,
  {
    id,
    userId,
    prompt,
    status,
    imageCount,
    durationMs,
    errorMessage,
    model,
    size,
    quality,
    outputFormat,
    background,
    moderation,
    outputCompression,
    createdAt,
  }: {
    id: string;
    userId: string;
    prompt: string;
    status: GenerationJobStatus;
    imageCount: number;
    durationMs: number | null;
    errorMessage: string | null;
    model: string | null;
    size: string | null;
    quality: string | null;
    outputFormat: string | null;
    background: string | null;
    moderation: string | null;
    outputCompression: number | null;
    createdAt: string;
  },
): void {
  database.prepare(
    `
    INSERT INTO generation_jobs (
      id,
      user_id,
      prompt,
      status,
      image_count,
      duration_ms,
      model,
      size,
      quality,
      output_format,
      background,
      moderation,
      output_compression,
      error_message,
      created_at
    ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
  `,
  ).run(
    id,
    userId,
    prompt,
    status,
    imageCount,
    durationMs,
    model,
    size,
    quality,
    outputFormat,
    background,
    moderation,
    outputCompression,
    errorMessage,
    createdAt,
  );
}

function insertGeneratedImages(
  database: SqliteDatabase,
  {
    jobId,
    userId,
    images,
    createdAt,
  }: {
    jobId: string;
    userId: string;
    images: StoredGeneratedImageInput[];
    createdAt: string;
  },
): void {
  const statement = database.prepare(
    `
    INSERT INTO generated_images (
      id,
      job_id,
      user_id,
      file_name,
      storage_key,
      mime_type,
      file_size,
      revised_prompt,
      created_at
    ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
  `,
  );

  for (const image of images) {
    statement.run(
      image.id,
      jobId,
      userId,
      image.fileName,
      image.storageKey,
      image.mimeType,
      image.fileSize,
      image.revisedPrompt,
      createdAt,
    );
  }
}

function mapImageHistoryItem(row: ImageHistoryRow): ImageHistoryItem {
  return {
    id: row.id,
    jobId: row.job_id,
    userId: row.user_id,
    username: row.username,
    prompt: row.prompt,
    revisedPrompt: row.revised_prompt,
    model: row.model,
    size: row.size,
    quality: row.quality,
    outputFormat: row.output_format,
    background: row.background,
    moderation: row.moderation,
    fileName: row.file_name,
    mimeType: row.mime_type,
    fileSize: Number(row.file_size),
    createdAt: row.created_at,
    imageUrl: buildStoredImageUrl(row.id),
    downloadUrl: buildStoredImageDownloadUrl(row.id),
  };
}

function mapUserRecord(row: UserRow | null | undefined): UserRecord | null {
  if (!row) {
    return null;
  }

  return {
    id: row.id,
    username: row.username,
    passwordHash: row.password_hash,
    role: row.role,
    monthlyQuota:
      row.monthly_quota === null || row.monthly_quota === undefined
        ? null
        : Number(row.monthly_quota),
    isActive: Boolean(row.is_active),
    createdAt: row.created_at,
    updatedAt: row.updated_at,
  };
}

function normalizeMonthlyQuota(value: unknown): number | null {
  if (value === null || value === undefined || value === "") {
    return null;
  }

  return normalizeQuotaInput(value);
}

function normalizeQuotaInput(value: unknown): number {
  const numericValue = Number.parseInt(String(value ?? ""), 10);

  if (!Number.isFinite(numericValue) || numericValue < 0) {
    return 0;
  }

  return numericValue;
}