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; 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; }