aigen/server/utils/library.ts

import { existsSync, mkdirSync, readFileSync, renameSync, rmSync, writeFileSync } from 'node:fs'
import { join } from 'node:path'
import { randomBytes, scrypt, timingSafeEqual } from 'node:crypto'
import { promisify } from 'node:util'
import type { H3Event } from 'h3'

const scryptAsync = promisify(scrypt)

export interface LibraryFolder {
  id: string
  name: string
  createdAt: number
}

export interface LibraryClip {
  id: string
  folderId: string
  prompt: string
  aspect: string
  width: number
  height: number
  steps: number
  turbo: boolean
  seed: number
  hideThumbnail: boolean
  createdAt: number
}

interface Catalog {
  passwordHash: string | null
  folders: LibraryFolder[]
  clips: LibraryClip[]
}

function libraryRoot() {
  const config = useRuntimeConfig()
  return (config.libraryDir || process.env.LIBRARY_DIR || '/data/library').replace(/\/$/, '')
}

function catalogPath() {
  return join(libraryRoot(), 'catalog.json')
}

function ensureRoot() {
  const root = libraryRoot()
  mkdirSync(join(root, 'files'), { recursive: true })
  return root
}

function emptyCatalog(): Catalog {
  return {
    passwordHash: null,
    folders: [{ id: 'inbox', name: 'Inbox', createdAt: Date.now() }],
    clips: []
  }
}

let writeChain = Promise.resolve()

function readCatalog(): Catalog {
  ensureRoot()
  if (!existsSync(catalogPath())) {
    const created = emptyCatalog()
    writeFileSync(catalogPath(), JSON.stringify(created, null, 2))
    return created
  }
  try {
    const parsed = JSON.parse(readFileSync(catalogPath(), 'utf8')) as Catalog
    if (!Array.isArray(parsed.folders) || !parsed.folders.length) {
      parsed.folders = emptyCatalog().folders
    }
    parsed.clips = Array.isArray(parsed.clips) ? parsed.clips : []
    parsed.passwordHash = parsed.passwordHash || null
    return parsed
  } catch {
    return emptyCatalog()
  }
}

function writeCatalog(catalog: Catalog) {
  ensureRoot()
  const tmp = catalogPath() + '.tmp'
  writeFileSync(tmp, JSON.stringify(catalog, null, 2))
  renameSync(tmp, catalogPath())
}

function mutate<T>(fn: (catalog: Catalog) => T): Promise<T> {
  const run = writeChain.then(() => {
    const catalog = readCatalog()
    const result = fn(catalog)
    writeCatalog(catalog)
    return result
  })
  writeChain = run.then(() => undefined, () => undefined)
  return run
}

export function passwordVersion(hash: string | null) {
  return hash ? hash.slice(-24) : 'open'
}

export function currentPasswordVersion() {
  return passwordVersion(readCatalog().passwordHash)
}

export function isLibraryLocked() {
  return Boolean(readCatalog().passwordHash)
}

export function libraryUnlocked(event: H3Event) {
  const catalog = readCatalog()
  if (!catalog.passwordHash) return true
  return getLibraryUnlockVersion(event) === passwordVersion(catalog.passwordHash)
}

export function assertLibraryAccess(event: H3Event) {
  if (libraryUnlocked(event)) return readCatalog()
  throw createError({ statusCode: 403, statusMessage: 'Library is locked' })
}

export async function hashLibraryPassword(password: string) {
  const salt = randomBytes(16)
  const hash = await scryptAsync(password, salt, 64) as Buffer
  return `scrypt:${salt.toString('hex')}:${hash.toString('hex')}`
}

export async function verifyLibraryPassword(password: string) {
  const stored = readCatalog().passwordHash
  if (!stored) return true
  const parts = stored.split(':')
  if (parts.length !== 3 || parts[0] !== 'scrypt') return false
  const salt = Buffer.from(parts[1], 'hex')
  const expected = Buffer.from(parts[2], 'hex')
  const actual = await scryptAsync(password, salt, 64) as Buffer
  return actual.length === expected.length && timingSafeEqual(actual, expected)
}

export function publicLibrary(event: H3Event) {
  const catalog = readCatalog()
  const locked = Boolean(catalog.passwordHash)
  const unlocked = libraryUnlocked(event)
  if (locked && !unlocked) {
    return { locked: true, unlocked: false, folders: [] as LibraryFolder[], clips: [] as LibraryClip[] }
  }
  return {
    locked,
    unlocked: true,
    folders: catalog.folders,
    clips: catalog.clips
  }
}

export function createFolder(name: string) {
  const trimmed = name.trim().slice(0, 80)
  if (!trimmed) throw createError({ statusCode: 400, statusMessage: 'Folder name is required' })
  return mutate((catalog) => {
    if (catalog.folders.some(folder => folder.name.toLowerCase() === trimmed.toLowerCase())) {
      throw createError({ statusCode: 409, statusMessage: 'A folder with that name already exists' })
    }
    const folder: LibraryFolder = { id: crypto.randomUUID(), name: trimmed, createdAt: Date.now() }
    catalog.folders.push(folder)
    return folder
  })
}

export function renameFolder(id: string, name: string) {
  const trimmed = name.trim().slice(0, 80)
  if (!trimmed) throw createError({ statusCode: 400, statusMessage: 'Folder name is required' })
  return mutate((catalog) => {
    const folder = catalog.folders.find(item => item.id === id)
    if (!folder) throw createError({ statusCode: 404, statusMessage: 'Folder not found' })
    folder.name = trimmed
    return folder
  })
}

export function deleteFolder(id: string) {
  return mutate((catalog) => {
    if (catalog.folders.length <= 1) {
      throw createError({ statusCode: 400, statusMessage: 'Keep at least one library folder' })
    }
    const folder = catalog.folders.find(item => item.id === id)
    if (!folder) throw createError({ statusCode: 404, statusMessage: 'Folder not found' })
    const clips = catalog.clips.filter(clip => clip.folderId === id)
    catalog.folders = catalog.folders.filter(item => item.id !== id)
    catalog.clips = catalog.clips.filter(clip => clip.folderId !== id)
    const fallback = catalog.folders[0].id
    for (const clip of clips) {
      rmSync(clipDir(clip.id), { recursive: true, force: true })
    }
    return { fallbackFolderId: fallback }
  })
}

export async function setLibraryPassword(password: string | null) {
  const hash = password && password.length ? await hashLibraryPassword(password) : null
  return mutate((catalog) => {
    catalog.passwordHash = hash
    return passwordVersion(hash)
  })
}

function clipDir(id: string) {
  return join(libraryRoot(), 'files', id)
}

export function clipVideoPath(id: string) {
  return join(clipDir(id), 'video.mp4')
}

export function clipThumbPath(id: string) {
  return join(clipDir(id), 'thumb')
}

export async function saveClip(params: {
  folderId: string
  prompt: string
  aspect: string
  width: number
  height: number
  steps: number
  turbo: boolean
  seed: number
  hideThumbnail: boolean
  video: Buffer
  thumb?: Buffer | null
}) {
  const catalog = readCatalog()
  const folder = catalog.folders.find(item => item.id === params.folderId) || catalog.folders[0]
  if (!folder) throw createError({ statusCode: 400, statusMessage: 'No library folder available' })
  const clip: LibraryClip = {
    id: crypto.randomUUID(),
    folderId: folder.id,
    prompt: params.prompt,
    aspect: params.aspect,
    width: params.width,
    height: params.height,
    steps: params.steps,
    turbo: params.turbo,
    seed: params.seed,
    hideThumbnail: params.hideThumbnail,
    createdAt: Date.now()
  }
  mkdirSync(clipDir(clip.id), { recursive: true })
  writeFileSync(clipVideoPath(clip.id), params.video)
  if (!params.hideThumbnail && params.thumb?.length) {
    writeFileSync(clipThumbPath(clip.id), params.thumb)
  }
  await mutate((next) => {
    next.clips.unshift(clip)
  })
  return clip
}

export function getClip(id: string) {
  const clip = readCatalog().clips.find(item => item.id === id)
  if (!clip) throw createError({ statusCode: 404, statusMessage: 'Clip not found' })
  return clip
}

export function deleteClip(id: string) {
  return mutate((catalog) => {
    const clip = catalog.clips.find(item => item.id === id)
    if (!clip) throw createError({ statusCode: 404, statusMessage: 'Clip not found' })
    catalog.clips = catalog.clips.filter(item => item.id !== id)
    rmSync(clipDir(id), { recursive: true, force: true })
    return clip
  })
}

export async function downloadComfyVideo(video: { filename: string; subfolder: string; type: string }) {
  const params = new URLSearchParams({
    filename: video.filename,
    subfolder: video.subfolder,
    type: video.type
  })
  const res = await comfyFetch(`/view?${params.toString()}`)
  if (!res.ok) throw new Error('Failed to fetch completed video from ComfyUI')
  return Buffer.from(await res.arrayBuffer())
}