aigen/server/utils/generationLog.ts

import { existsSync, mkdirSync, readFileSync, renameSync, writeFileSync } from 'node:fs'
import { join } from 'node:path'
import { ownerRoot } from '~/server/utils/library'
import type { StudioJobKind, StudioJobPayload } from '~/server/utils/studioQueue'

export type GenerationLogStatus = 'queued' | 'running' | 'complete' | 'error' | 'cancelled'

export type GenerationLogEntry = {
  id: string
  studioJobId?: string
  createdAt: number
  updatedAt: number
  kind: StudioJobKind
  name: string
  prompt: string
  status: GenerationLogStatus
  lastError?: string
  payload: StudioJobPayload
}

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

function logPath(owner: string) {
  return join(ownerRoot(owner), 'generation-log.json')
}

function ensureOwner(owner: string) {
  mkdirSync(ownerRoot(owner), { recursive: true })
}

function readLog(owner: string): GenerationLogEntry[] {
  ensureOwner(owner)
  const path = logPath(owner)
  if (!existsSync(path)) return []
  try {
    const parsed = JSON.parse(readFileSync(path, 'utf8'))
    return Array.isArray(parsed) ? parsed.filter(item => item && typeof item === 'object' && item.id && item.payload) : []
  } catch {
    return []
  }
}

function writeLog(owner: string, entries: GenerationLogEntry[]) {
  ensureOwner(owner)
  const path = logPath(owner)
  const tmp = `${path}.tmp`
  writeFileSync(tmp, JSON.stringify(entries, null, 2))
  renameSync(tmp, path)
}

function mutate<T>(owner: string, fn: (entries: GenerationLogEntry[]) => T): Promise<T> {
  const prev = writeChains.get(owner) || Promise.resolve()
  const run = prev.then(() => {
    const entries = readLog(owner)
    const result = fn(entries)
    writeLog(owner, entries)
    return result
  })
  writeChains.set(owner, run.then(() => undefined, () => undefined))
  return run
}

function snippetName(payload: StudioJobPayload, kind: StudioJobKind) {
  const named = String(payload.name || '').trim()
  if (named) return named.slice(0, 80)
  if (kind === 'music') return String(payload.prompt || 'Track').slice(0, 80)
  return String(payload.prompt || 'Untitled').slice(0, 80)
}

export function listGenerationLog(owner: string, limit = 40) {
  return readLog(owner).slice(0, Math.max(1, Math.min(MAX_ENTRIES, limit)))
}

export function getGenerationLogEntry(owner: string, id: string) {
  const entry = readLog(owner).find(item => item.id === id)
  if (!entry) throw createError({ statusCode: 404, statusMessage: 'Generation log entry not found' })
  return entry
}

export async function recordGenerationQueued(params: {
  ownerKey: string
  studioJobId: string
  kind: StudioJobKind
  payload: StudioJobPayload
}) {
  const now = Date.now()
  return mutate(params.ownerKey, (entries) => {
    const existing = entries.find(item => item.studioJobId === params.studioJobId)
    if (existing) {
      existing.updatedAt = now
      existing.status = 'queued'
      existing.payload = structuredClone(params.payload)
      existing.name = snippetName(params.payload, params.kind)
      existing.prompt = params.payload.prompt || existing.prompt
      existing.kind = params.kind
      existing.lastError = undefined
      entries.splice(entries.indexOf(existing), 1)
      entries.unshift(existing)
      return existing
    }
    const entry: GenerationLogEntry = {
      id: crypto.randomUUID(),
      studioJobId: params.studioJobId,
      createdAt: now,
      updatedAt: now,
      kind: params.kind,
      name: snippetName(params.payload, params.kind),
      prompt: params.payload.prompt || '',
      status: 'queued',
      payload: structuredClone(params.payload)
    }
    entries.unshift(entry)
    entries.splice(MAX_ENTRIES)
    return entry
  })
}

export async function updateGenerationLogByStudioJob(
  owner: string,
  studioJobId: string,
  patch: { status?: GenerationLogStatus; lastError?: string; payload?: StudioJobPayload }
) {
  return mutate(owner, (entries) => {
    const entry = entries.find(item => item.studioJobId === studioJobId)
    if (!entry) return null
    entry.updatedAt = Date.now()
    if (patch.status) entry.status = patch.status
    if (patch.lastError !== undefined) entry.lastError = patch.lastError || undefined
    if (patch.payload) {
      entry.payload = structuredClone(patch.payload)
      entry.name = snippetName(patch.payload, entry.kind)
      entry.prompt = patch.payload.prompt || entry.prompt
    }
    return structuredClone(entry)
  })
}

export async function requeueGenerationLog(owner: string, id: string) {
  const entry = getGenerationLogEntry(owner, id)
  const { addStudioJob, kickStudioQueue } = await import('~/server/utils/studioQueue')
  let payload = structuredClone(entry.payload)
  if (payload?.studio2Id || payload?.studio2Request) {
    const { saveRecord, readRecord } = await import('~/server/utils/studio2/store')
    let request = payload.studio2Request
    if (!request && payload.studio2Id) {
      try { request = structuredClone(readRecord(owner, payload.studio2Id).request) }
      catch { request = null }
    }
    if (!request) throw createError({ statusCode: 400, statusMessage: 'That Studio 2 job has no saved request to re-queue.' })
    const studio2Id = crypto.randomUUID()
    saveRecord({
      id: studio2Id,
      owner,
      request: structuredClone(request),
      state: 'waiting',
      queuedAt: Date.now(),
      familyId: crypto.randomUUID(),
      outputs: [],
      purgeResult: 'Not yet saved'
    })
    payload = {
      ...payload,
      ...structuredClone(request.settings || {}),
      studio2Id,
      studio2Request: structuredClone(request),
      prompt: request.compiledPrompt || payload.prompt,
      folderId: request.folderId || payload.folderId,
      extensions: [],
      referenceStillIds: [],
      useIdentityRefs: false,
      queueAutoRun: true
    }
  }
  const studio = await addStudioJob({
    ownerKey: owner,
    familyId: crypto.randomUUID(),
    kind: entry.kind,
    payload
  })
  await kickStudioQueue()
  return { studioJobId: studio.id, entry: listGenerationLog(owner, 1)[0] || getGenerationLogEntry(owner, id) }
}

export async function deleteGenerationLogEntry(owner: string, id: string) {
  return mutate(owner, (entries) => {
    const index = entries.findIndex(item => item.id === id)
    if (index < 0) throw createError({ statusCode: 404, statusMessage: 'Generation log entry not found' })
    const [removed] = entries.splice(index, 1)
    return removed
  })
}