aigen/server/utils/library.ts

import { createHash } from 'node:crypto'
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
  thumbnailDefault: boolean
  passwordHash: string | null
}

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
  comfyFilename?: string
}

export interface LibraryStill {
  id: string
  folderId: string
  filename: string
  width: number
  height: number
  hash: string
  createdAt: number
}

export interface PublicFolder {
  id: string
  name: string
  createdAt: number
  thumbnailDefault: boolean
  protected: boolean
  unlocked: boolean
}

interface Catalog {
  folders: LibraryFolder[]
  clips: LibraryClip[]
  stills: LibraryStill[]
}

const writeChains = new Map<string, Promise<unknown>>()

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

export function libraryOwnerKey(event: H3Event) {
  const user = getSessionUser(event)
  const raw = user?.sub || user?.email || (!authEnabled() ? 'local' : '')
  if (!raw) {
    throw createError({ statusCode: 401, statusMessage: 'Authentication required' })
  }
  return createHash('sha256').update(raw).digest('hex').slice(0, 24)
}

function ownerRoot(owner: string) {
  return join(libraryRoot(), 'users', owner)
}

function catalogPath(owner: string) {
  return join(ownerRoot(owner), 'catalog.json')
}

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

function normalizeFolder(folder: LibraryFolder): LibraryFolder {
  return {
    id: folder.id,
    name: folder.name,
    createdAt: folder.createdAt || Date.now(),
    thumbnailDefault: folder.thumbnailDefault !== false,
    passwordHash: folder.passwordHash || null
  }
}

function normalizeCatalog(parsed: Partial<Catalog> & { passwordHash?: string | null }): Catalog {
  const folders = Array.isArray(parsed.folders) && parsed.folders.length
    ? parsed.folders.map(normalizeFolder)
    : emptyCatalog().folders
  if (parsed.passwordHash) {
    for (const folder of folders) {
      if (!folder.passwordHash) folder.passwordHash = parsed.passwordHash
    }
  }
  return {
    folders,
    clips: Array.isArray(parsed.clips) ? parsed.clips : [],
    stills: Array.isArray(parsed.stills) ? parsed.stills : []
  }
}

function migrateLegacy(owner: string) {
  const dest = catalogPath(owner)
  if (existsSync(dest)) return
  mkdirSync(ownerRoot(owner), { recursive: true })
  const legacyCatalog = join(libraryRoot(), 'catalog.json')
  const legacyFiles = join(libraryRoot(), 'files')
  if (!existsSync(legacyCatalog)) {
    writeFileSync(dest, JSON.stringify(emptyCatalog(), null, 2))
    return
  }
  try {
    const catalog = normalizeCatalog(JSON.parse(readFileSync(legacyCatalog, 'utf8')))
    writeFileSync(dest, JSON.stringify(catalog, null, 2))
    if (existsSync(legacyFiles) && !existsSync(join(ownerRoot(owner), 'files'))) {
      renameSync(legacyFiles, join(ownerRoot(owner), 'files'))
    }
    renameSync(legacyCatalog, `${legacyCatalog}.migrated`)
  } catch {
    writeFileSync(dest, JSON.stringify(emptyCatalog(), null, 2))
  }
}

function ensureOwner(owner: string) {
  mkdirSync(join(ownerRoot(owner), 'files'), { recursive: true })
  mkdirSync(join(ownerRoot(owner), 'stills'), { recursive: true })
  migrateLegacy(owner)
}

function readCatalog(owner: string): Catalog {
  ensureOwner(owner)
  try {
    return normalizeCatalog(JSON.parse(readFileSync(catalogPath(owner), 'utf8')))
  } catch {
    const created = emptyCatalog()
    writeFileSync(catalogPath(owner), JSON.stringify(created, null, 2))
    return created
  }
}

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

function mutate<T>(owner: string, fn: (catalog: Catalog) => T): Promise<T> {
  const prev = writeChains.get(owner) || Promise.resolve()
  const run = prev.then(() => {
    const catalog = readCatalog(owner)
    const result = fn(catalog)
    writeCatalog(owner, catalog)
    return result
  })
  writeChains.set(owner, run.then(() => undefined, () => undefined))
  return run
}

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

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')}`
}

async function verifyStoredPassword(stored: string | null, password: string) {
  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)
}

function folderUnlocks(event: H3Event) {
  return getFolderUnlocks(event)
}

export function folderUnlocked(event: H3Event, folder: LibraryFolder) {
  if (!folder.passwordHash) return true
  return folderUnlocks(event)[folder.id] === passwordVersion(folder.passwordHash)
}

export function publicFolder(event: H3Event, folder: LibraryFolder): PublicFolder {
  return {
    id: folder.id,
    name: folder.name,
    createdAt: folder.createdAt,
    thumbnailDefault: folder.thumbnailDefault,
    protected: Boolean(folder.passwordHash),
    unlocked: folderUnlocked(event, folder)
  }
}

export function publicLibrary(event: H3Event) {
  const owner = libraryOwnerKey(event)
  const catalog = readCatalog(owner)
  const folders = catalog.folders.map(folder => publicFolder(event, folder))
  const unlockedIds = new Set(folders.filter(folder => folder.unlocked).map(folder => folder.id))
  return {
    folders,
    clips: catalog.clips.filter(clip => unlockedIds.has(clip.folderId)),
    stills: catalog.stills.filter(still => unlockedIds.has(still.folderId))
  }
}

export function assertLibraryOwner(event: H3Event) {
  const owner = libraryOwnerKey(event)
  return { owner, catalog: readCatalog(owner) }
}

export function assertFolderAccess(event: H3Event, folderId: string) {
  const { owner, catalog } = assertLibraryOwner(event)
  const folder = catalog.folders.find(item => item.id === folderId)
  if (!folder) throw createError({ statusCode: 404, statusMessage: 'Folder not found' })
  if (!folderUnlocked(event, folder)) {
    throw createError({ statusCode: 403, statusMessage: 'Folder is locked' })
  }
  return { owner, catalog, folder }
}

export function assertLibraryAccess(event: H3Event) {
  return assertLibraryOwner(event).catalog
}

export function createFolder(owner: string, name: string) {
  const trimmed = name.trim().slice(0, 80)
  if (!trimmed) throw createError({ statusCode: 400, statusMessage: 'Folder name is required' })
  return mutate(owner, (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(),
      thumbnailDefault: true,
      passwordHash: null
    }
    catalog.folders.push(folder)
    return folder
  })
}

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

export function deleteFolder(owner: string, id: string) {
  return mutate(owner, (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)
    const stills = catalog.stills.filter(still => still.folderId === id)
    catalog.folders = catalog.folders.filter(item => item.id !== id)
    catalog.clips = catalog.clips.filter(clip => clip.folderId !== id)
    catalog.stills = catalog.stills.filter(still => still.folderId !== id)
    for (const clip of clips) {
      rmSync(clipDir(owner, clip.id), { recursive: true, force: true })
    }
    for (const still of stills) {
      rmSync(stillPath(owner, still.id), { force: true })
    }
    return { fallbackFolderId: catalog.folders[0].id }
  })
}

export async function updateFolder(owner: string, id: string, patch: {
  name?: string
  thumbnailDefault?: boolean
  password?: string | null
  currentPassword?: string
}) {
  const catalog = readCatalog(owner)
  const folder = catalog.folders.find(item => item.id === id)
  if (!folder) throw createError({ statusCode: 404, statusMessage: 'Folder not found' })

  let nextName = folder.name
  if (patch.name !== undefined) {
    nextName = patch.name.trim().slice(0, 80)
    if (!nextName) throw createError({ statusCode: 400, statusMessage: 'Folder name is required' })
    if (catalog.folders.some(item => item.id !== id && item.name.toLowerCase() === nextName.toLowerCase())) {
      throw createError({ statusCode: 409, statusMessage: 'A folder with that name already exists' })
    }
  }

  let nextHash: string | null | undefined
  if (patch.password !== undefined) {
    const next = String(patch.password || '').trim()
    if (folder.passwordHash && !(await verifyStoredPassword(folder.passwordHash, String(patch.currentPassword || '')))) {
      throw createError({ statusCode: 401, statusMessage: 'Current folder password is incorrect' })
    }
    if (next && next.length < 4) {
      throw createError({ statusCode: 400, statusMessage: 'Use at least 4 characters, or leave blank to remove the password' })
    }
    nextHash = next ? await hashLibraryPassword(next) : null
  }

  return mutate(owner, (nextCatalog) => {
    const item = nextCatalog.folders.find(entry => entry.id === id)
    if (!item) throw createError({ statusCode: 404, statusMessage: 'Folder not found' })
    item.name = nextName
    if (typeof patch.thumbnailDefault === 'boolean') item.thumbnailDefault = patch.thumbnailDefault
    if (nextHash !== undefined) item.passwordHash = nextHash
    return item
  })
}

export async function verifyFolderPassword(owner: string, folderId: string, password: string) {
  const folder = readCatalog(owner).folders.find(item => item.id === folderId)
  if (!folder) throw createError({ statusCode: 404, statusMessage: 'Folder not found' })
  if (!folder.passwordHash) return folder
  if (!(await verifyStoredPassword(folder.passwordHash, password))) {
    throw createError({ statusCode: 401, statusMessage: 'Incorrect folder password' })
  }
  return folder
}

export async function unlockMatchingFolders(owner: string, password: string, folderId?: string) {
  const catalog = readCatalog(owner)
  const targets = folderId
    ? catalog.folders.filter(folder => folder.id === folderId)
    : catalog.folders.filter(folder => folder.passwordHash)
  if (folderId && !targets.length) {
    throw createError({ statusCode: 404, statusMessage: 'Folder not found' })
  }
  const unlocked: LibraryFolder[] = []
  for (const folder of targets) {
    if (!folder.passwordHash || await verifyStoredPassword(folder.passwordHash, password)) {
      unlocked.push(folder)
    }
  }
  if (!unlocked.length) {
    throw createError({ statusCode: 401, statusMessage: 'Incorrect folder password' })
  }
  return unlocked
}

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

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

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

export function stillPath(owner: string, id: string) {
  return join(ownerRoot(owner), 'stills', id)
}

function bytesHash(data: Buffer) {
  return createHash('sha256').update(data).digest('hex')
}

export async function saveStill(params: {
  ownerKey: string
  folderId: string
  filename: string
  data: Buffer
  width?: number
  height?: number
}) {
  const catalog = readCatalog(params.ownerKey)
  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 hash = bytesHash(params.data)
  const existing = catalog.stills.find(item => item.folderId === folder.id && item.hash === hash)
  if (existing) return existing
  const still: LibraryStill = {
    id: crypto.randomUUID(),
    folderId: folder.id,
    filename: params.filename.slice(0, 120) || 'still.png',
    width: params.width || 0,
    height: params.height || 0,
    hash,
    createdAt: Date.now()
  }
  mkdirSync(join(ownerRoot(params.ownerKey), 'stills'), { recursive: true })
  writeFileSync(stillPath(params.ownerKey, still.id), params.data)
  await mutate(params.ownerKey, (next) => {
    next.stills.unshift(still)
  })
  return still
}

export function getStill(owner: string, id: string) {
  const still = readCatalog(owner).stills.find(item => item.id === id)
  if (!still) throw createError({ statusCode: 404, statusMessage: 'Image not found' })
  return still
}

export function deleteStill(owner: string, id: string) {
  return mutate(owner, (catalog) => {
    const still = catalog.stills.find(item => item.id === id)
    if (!still) throw createError({ statusCode: 404, statusMessage: 'Image not found' })
    catalog.stills = catalog.stills.filter(item => item.id !== id)
    rmSync(stillPath(owner, id), { force: true })
    return still
  })
}

export function imageContentType(buf: Buffer) {
  if (buf[0] === 0x89 && buf[1] === 0x50) return 'image/png'
  if (buf[0] === 0xff && buf[1] === 0xd8) return 'image/jpeg'
  if (buf[0] === 0x52 && buf[1] === 0x49) return 'image/webp'
  return 'application/octet-stream'
}

export async function saveClip(params: {
  ownerKey: string
  folderId: string
  prompt: string
  aspect: string
  width: number
  height: number
  steps: number
  turbo: boolean
  seed: number
  hideThumbnail: boolean
  video: Buffer
  thumb?: Buffer | null
  comfyFilename?: string
}) {
  const catalog = readCatalog(params.ownerKey)
  const folder = catalog.folders.find(item => item.id === params.folderId) || catalog.folders[0]
  if (!folder) throw createError({ statusCode: 400, statusMessage: 'No library folder available' })
  if (params.comfyFilename && catalog.clips.some(clip => clip.comfyFilename === params.comfyFilename)) {
    return catalog.clips.find(clip => clip.comfyFilename === params.comfyFilename) as LibraryClip
  }
  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(),
    comfyFilename: params.comfyFilename
  }
  mkdirSync(clipDir(params.ownerKey, clip.id), { recursive: true })
  writeFileSync(clipVideoPath(params.ownerKey, clip.id), params.video)
  if (!params.hideThumbnail && params.thumb?.length) {
    writeFileSync(clipThumbPath(params.ownerKey, clip.id), params.thumb)
  }
  if (params.thumb?.length) {
    await saveStill({
      ownerKey: params.ownerKey,
      folderId: folder.id,
      filename: `${clip.id}.png`,
      data: params.thumb,
      width: params.width,
      height: params.height
    })
  }
  await mutate(params.ownerKey, (next) => {
    next.clips.unshift(clip)
  })
  return clip
}

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

export function deleteClip(owner: string, id: string) {
  return mutate(owner, (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(owner, id), { recursive: true, force: true })
    return clip
  })
}

export async function downloadComfyVideo(video: { filename: string; subfolder: string; type: string }) {
  const subfolders = [...new Set([video.subfolder, 'video', ''])]
  let lastStatus = 0
  for (const subfolder of subfolders) {
    const params = new URLSearchParams({
      filename: video.filename,
      subfolder,
      type: video.type || 'output'
    })
    const res = await comfyFetch(`/view?${params.toString()}`)
    lastStatus = res.status
    if (res.ok) return Buffer.from(await res.arrayBuffer())
  }
  throw new Error(`Failed to fetch completed video from ComfyUI (${lastStatus})`)
}

export async function importMissingComfyVideos(owner: string, folderId?: string) {
  const catalog = readCatalog(owner)
  const folder = catalog.folders.find(item => item.id === folderId) || catalog.folders[0]
  if (!folder) return []
  const known = new Set(catalog.clips.map(clip => clip.comfyFilename).filter(Boolean) as string[])
  const unlabeled = catalog.clips.filter(clip => !clip.comfyFilename).length
  const history = await fetchHistoryAll()
  const found: { promptId: string; video: { filename: string; subfolder: string; type: string }; prompt: string }[] = []
  for (const [promptId, entry] of Object.entries(history)) {
    const video = extractVideo({ [promptId]: entry as Record<string, unknown> }, promptId)
    if (!video || known.has(video.filename)) continue
    found.push({
      promptId,
      video,
      prompt: extractPromptFromHistory(entry) || 'Recovered from ComfyUI'
    })
  }
  found.sort((a, b) => a.video.filename.localeCompare(b.video.filename, undefined, { numeric: true }))
  const toImport = found.slice(unlabeled)
  const imported = []
  for (const item of toImport) {
    const buffer = await downloadComfyVideo(item.video)
    const clip = await saveClip({
      ownerKey: owner,
      folderId: folder.id,
      prompt: item.prompt,
      aspect: 'auto',
      width: 0,
      height: 0,
      steps: 0,
      turbo: true,
      seed: 0,
      hideThumbnail: false,
      video: buffer,
      thumb: null,
      comfyFilename: item.video.filename
    })
    imported.push(clip)
  }
  return imported
}