aigen/server/utils/shotQueue.ts

import { existsSync, mkdirSync, readdirSync, readFileSync, renameSync, writeFileSync } from 'node:fs'
import { join } from 'node:path'
import { getJob, listJobs, type Job } from '~/server/utils/jobs'
import { parsePostedLoraStack, persistLoraFields } from '~/server/utils/loras'
import type { LoraStackItem } from '~/utils/loras'
import { persistPromptWrappers } from '~/utils/promptParts'
import type { QueuedExtension } from '~/server/utils/library'

export type ShotQueueStatus = 'idle' | 'running' | 'paused' | 'complete' | 'error'
export type ShotSegmentStatus = 'pending' | 'running' | 'complete' | 'error'

export interface ShotQueueSegment {
  index: number
  prompt: string
  promptPre?: string
  promptPost?: string
  duration: number
  status: ShotSegmentStatus
  clipId?: string
  error?: string
  permanenceRefs?: import('~/utils/globalLocks').PermanenceRef[]
  loraName?: string
  loraStack?: LoraStackItem[]
}

export interface ShotQueue {
  id: string
  name: string
  familyId: string
  folderId: string
  createdAt: number
  updatedAt: number
  status: ShotQueueStatus
  autoRun: boolean
  stopAfterCurrent?: boolean
  currentClipId?: string
  currentJobId?: string
  stillId?: string
  stillFilename?: string
  hideThumbnail: boolean
  hideInput?: boolean
  aspect: string
  width: number
  height: number
  steps: number
  turbo: boolean
  cfg?: number
  fps?: number
  samplerName?: string
  scheduler?: string
  workflow?: import('~/utils/videoModels').VideoWorkflowId
  sound?: boolean
  useIdentityRefs?: boolean
  referenceStillIds?: Array<string | null>
  globalLocks?: string
  permanenceRefs?: import('~/utils/globalLocks').PermanenceRef[]
  loraName?: string
  loraStack?: LoraStackItem[]
  promptPre?: string
  promptPost?: string
  segments: ShotQueueSegment[]
  lastError?: string
}

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

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

function queuesPath(owner: string) {
  return join(libraryRoot(), 'users', owner, 'shot-queues.json')
}

function ensureOwner(owner: string) {
  mkdirSync(join(libraryRoot(), 'users', owner), { recursive: true })
}

function readQueues(owner: string): ShotQueue[] {
  ensureOwner(owner)
  const path = queuesPath(owner)
  if (!existsSync(path)) return []
  try {
    const parsed = JSON.parse(readFileSync(path, 'utf8'))
    return Array.isArray(parsed) ? parsed : []
  } catch {
    return []
  }
}

function writeQueues(owner: string, queues: ShotQueue[]) {
  ensureOwner(owner)
  const path = queuesPath(owner)
  const tmp = `${path}.tmp`
  writeFileSync(tmp, JSON.stringify(queues, null, 2))
  renameSync(tmp, path)
}

function mutate<T>(owner: string, fn: (queues: ShotQueue[]) => T): Promise<T> {
  const prev = writeChains.get(owner) || Promise.resolve()
  const run = prev.then(() => {
    const queues = readQueues(owner)
    const result = fn(queues)
    writeQueues(owner, queues)
    return result
  })
  writeChains.set(owner, run.then(() => undefined, () => undefined))
  return run
}

export function pendingSegmentCount(queue: ShotQueue) {
  return queue.segments.filter(segment => segment.status === 'pending' || segment.status === 'error').length
}

export function completedSegmentCount(queue: ShotQueue) {
  return queue.segments.filter(segment => segment.status === 'complete').length
}

export function summarizeQueue(queue: ShotQueue) {
  return {
    ...queue,
    pendingCount: pendingSegmentCount(queue),
    completedCount: completedSegmentCount(queue),
    totalCount: queue.segments.length
  }
}

export function listShotQueues(owner: string) {
  return readQueues(owner)
    .slice()
    .sort((a, b) => b.updatedAt - a.updatedAt)
}

export function getShotQueue(owner: string, id: string) {
  return readQueues(owner).find(queue => queue.id === id) || null
}

export function queueIsProcessing(id: string) {
  return activeJobs.has(id)
}

export function markShotQueueCancelled(id: string) {
  if (!id) return
  cancelledQueueIds.add(id)
  activeJobs.delete(id)
}

export function isShotQueueCancelled(id: string) {
  return Boolean(id) && cancelledQueueIds.has(id)
}

export function setQueueJob(id: string, jobId: string | null) {
  if (jobId) {
    if (cancelledQueueIds.has(id)) return
    activeJobs.set(id, jobId)
  } else {
    activeJobs.delete(id)
  }
}

export function getQueueJobId(id: string) {
  return activeJobs.get(id)
}

export async function createShotQueue(params: {
  ownerKey: string
  name?: string
  familyId: string
  folderId: string
  autoRun: boolean
  stillId?: string
  stillFilename?: string
  hideThumbnail: boolean
  hideInput?: boolean
  aspect: string
  width: number
  height: number
  steps: number
  turbo: boolean
  cfg?: number
  fps?: number
  samplerName?: string
  scheduler?: string
  workflow?: import('~/utils/videoModels').VideoWorkflowId
  sound?: boolean
  useIdentityRefs?: boolean
  referenceStillIds?: Array<string | null>
  globalLocks?: string
  permanenceRefs?: import('~/utils/globalLocks').PermanenceRef[]
  loraName?: string
  loraStack?: LoraStackItem[]
  promptPre?: string
  promptPost?: string
  initial: { prompt: string; promptPre?: string; promptPost?: string; duration: number; permanenceRefs?: import('~/utils/globalLocks').PermanenceRef[]; loraName?: string; loraStack?: LoraStackItem[] }
  extensions: QueuedExtension[]
  jobId?: string
}): Promise<ShotQueue> {
  const now = Date.now()
  const queue: ShotQueue = {
    id: crypto.randomUUID(),
    name: (params.name || '').trim() || params.initial.prompt.slice(0, 80),
    familyId: params.familyId,
    folderId: params.folderId,
    createdAt: now,
    updatedAt: now,
    status: 'running',
    autoRun: params.autoRun,
    stopAfterCurrent: false,
    currentJobId: params.jobId,
    stillId: params.stillId,
    stillFilename: params.stillFilename,
    hideThumbnail: params.hideThumbnail,
    hideInput: params.hideInput,
    aspect: params.aspect,
    width: params.width,
    height: params.height,
    steps: params.steps,
    turbo: params.turbo,
    cfg: params.cfg,
    fps: params.fps,
    samplerName: params.samplerName,
    scheduler: params.scheduler,
    workflow: params.workflow,
    sound: params.sound,
    useIdentityRefs: params.useIdentityRefs,
    referenceStillIds: params.referenceStillIds,
    globalLocks: params.globalLocks,
    permanenceRefs: params.permanenceRefs,
    loraName: params.loraName,
    loraStack: params.loraStack,
    ...persistPromptWrappers(params),
    segments: [
      {
        index: 0,
        prompt: params.initial.prompt,
        duration: params.initial.duration,
        status: 'running',
        permanenceRefs: params.initial.permanenceRefs,
        ...persistLoraFields(params.initial.loraStack || params.initial.loraName),
        ...persistPromptWrappers(params),
        ...persistPromptWrappers(params.initial)
      },
      ...params.extensions.map((item, index) => ({
        index: index + 1,
        prompt: item.prompt,
        duration: item.duration,
        status: 'pending' as const,
        permanenceRefs: item.permanenceRefs,
        ...persistLoraFields(item.loraStack || item.loraName),
        ...persistPromptWrappers(params),
        ...persistPromptWrappers(item)
      }))
    ]
  }
  await mutate(params.ownerKey, (queues) => {
    queues.unshift(queue)
    return queue
  })
  if (params.jobId) setQueueJob(queue.id, params.jobId)
  return queue
}

export async function updateShotQueue(owner: string, id: string, patch: (queue: ShotQueue) => void) {
  if (cancelledQueueIds.has(id)) {
    throw createError({ statusCode: 410, statusMessage: 'This batch was removed' })
  }
  return mutate(owner, (queues) => {
    const queue = queues.find(item => item.id === id)
    if (!queue) throw createError({ statusCode: 404, statusMessage: 'Queue not found' })
    patch(queue)
    queue.updatedAt = Date.now()
    return structuredClone(queue)
  })
}

export async function patchShotQueueSegments(
  owner: string,
  id: string,
  segments: Array<{ index: number; prompt?: string; duration?: number; loraName?: string; loraStack?: unknown }>
) {
  return updateShotQueue(owner, id, (queue) => {
    if (queue.status === 'complete') {
      throw createError({ statusCode: 409, statusMessage: 'This batch is already finished' })
    }
    for (const next of segments) {
      const segment = queue.segments.find(item => item.index === next.index)
      if (!segment) continue
      if (segment.status === 'complete' || segment.status === 'running') continue
      if (typeof next.prompt === 'string') {
        const prompt = next.prompt.trim()
        if (!prompt) throw createError({ statusCode: 400, statusMessage: `Shot ${next.index + 1} needs a prompt` })
        segment.prompt = prompt
      }
      if (typeof next.duration === 'number' && Number.isFinite(next.duration)) {
        segment.duration = Math.min(120, Math.max(0.5, Math.round(next.duration * 10) / 10))
      }
      if ('loraStack' in next || 'loraName' in next) {
        const stack = parsePostedLoraStack(next.loraStack ?? next.loraName, 'video')
        const fields = persistLoraFields(stack)
        if (fields.loraStack) {
          segment.loraStack = fields.loraStack
          segment.loraName = fields.loraName
        } else {
          delete segment.loraStack
          delete segment.loraName
        }
      }
    }
  })
}

export async function markQueueSegment(owner: string, id: string, index: number, status: ShotSegmentStatus, extra: Partial<ShotQueueSegment> = {}) {
  return updateShotQueue(owner, id, (queue) => {
    const segment = queue.segments.find(item => item.index === index)
    if (!segment) return
    segment.status = status
    if (extra.clipId) segment.clipId = extra.clipId
    if (extra.error) segment.error = extra.error
    else if (status !== 'error') delete segment.error
    if (extra.clipId) queue.currentClipId = extra.clipId
    if (status === 'complete' && pendingSegmentCount(queue) === 0) {
      queue.status = 'complete'
      queue.autoRun = false
      queue.currentJobId = undefined
    }
  })
}

export async function syncQueueFromJob(job: Job, clipId?: string) {
  const library = job.library
  if (!library?.queueId) return null
  const index = library.chainIndex || 0
  const failed = job.status === 'error' || job.status === 'cancelled'
  if (failed) {
    return updateShotQueue(library.ownerKey, library.queueId, (queue) => {
      const segment = queue.segments.find(item => item.index === index)
      if (segment && segment.status !== 'complete') {
        segment.status = 'error'
        segment.error = job.error || 'Generation failed'
      }
      queue.status = job.status === 'cancelled' ? 'paused' : 'error'
      queue.lastError = job.error
      queue.currentJobId = undefined
      queue.autoRun = job.status === 'cancelled' ? false : queue.autoRun
    }).catch(() => null)
  }
  if (!clipId) return getShotQueue(library.ownerKey, library.queueId)
  return markQueueSegment(library.ownerKey, library.queueId, index, 'complete', { clipId })
}

export async function prepareQueueBurst(owner: string, id: string, count: number | 'all') {
  if (cancelledQueueIds.has(id)) {
    throw createError({ statusCode: 410, statusMessage: 'This batch was removed' })
  }
  return mutate(owner, (queues) => {
    const queue = queues.find(item => item.id === id)
    if (!queue) throw createError({ statusCode: 404, statusMessage: 'Queue not found' })
    if (activeJobs.has(queue.id)) {
      throw createError({ statusCode: 409, statusMessage: 'This batch is already processing' })
    }
    const remaining = queue.segments.filter(segment => (
      segment.status === 'pending'
      || segment.status === 'error'
      || (segment.status === 'running' && !activeJobs.has(queue.id))
    ))
    if (!remaining.length) {
      throw createError({ statusCode: 400, statusMessage: 'No remaining shots to process' })
    }
    const initial = queue.segments[0]
    if (initial && initial.status !== 'complete') {
      throw createError({ statusCode: 409, statusMessage: 'The first shot is still running or did not save. Wait for it, or generate again.' })
    }
    if (!queue.currentClipId) {
      throw createError({ statusCode: 409, statusMessage: 'The previous clip is missing, so the next shot cannot start' })
    }
    const n = count === 'all' ? remaining.length : Math.min(remaining.length, Math.max(1, Math.floor(count)))
    remaining.slice(0, n).forEach((segment) => {
      segment.status = segment.index === remaining[0].index ? 'running' : 'pending'
      delete segment.error
    })
    queue.autoRun = count === 'all'
    queue.stopAfterCurrent = false
    queue.status = 'running'
    queue.lastError = undefined
    queue.updatedAt = Date.now()
    return { queue: structuredClone(queue), count: n }
  })
}

export async function finishQueueBurst(owner: string, id: string, paused: boolean, error?: string) {
  return updateShotQueue(owner, id, (queue) => {
    queue.currentJobId = undefined
    for (const segment of queue.segments) {
      if (segment.status !== 'running') continue
      if (error) {
        segment.status = 'error'
        segment.error = error
      } else {
        segment.status = 'pending'
        delete segment.error
      }
    }
    if (error) {
      queue.status = 'error'
      queue.lastError = error
      queue.autoRun = false
      return
    }
    if (pendingSegmentCount(queue) === 0) {
      queue.status = 'complete'
      queue.autoRun = false
      return
    }
    queue.status = paused || !queue.autoRun ? 'paused' : 'idle'
    if (paused) queue.autoRun = false
  })
}

function liveJobForQueue(id: string): Job | undefined {
  const mapped = activeJobs.get(id)
  if (mapped) {
    const job = getJob(mapped)
    if (job) return job
  }
  return listJobs().find(job => (
    job.library?.queueId === id
    && (job.status === 'queued' || job.status === 'uploading' || job.status === 'running')
  ))
}

export async function setShotQueuePause(owner: string, id: string, pause: boolean) {
  const live = liveJobForQueue(id)
  if (live?.library) {
    live.library.stopAfterCurrent = pause
    if (pause) live.library.queueAutoRun = false
  }
  return updateShotQueue(owner, id, (queue) => {
    if (pause) {
      queue.autoRun = false
      queue.stopAfterCurrent = true
      if (queue.status === 'running') queue.lastError = undefined
      else if (queue.status !== 'complete') queue.status = 'paused'
    } else {
      queue.stopAfterCurrent = false
      if (live?.library) live.library.queueAutoRun = queue.autoRun === true
      if (queue.status === 'paused') queue.status = 'idle'
    }
  })
}

export async function pauseShotQueue(owner: string, id: string) {
  return setShotQueuePause(owner, id, true)
}

export async function deleteShotQueue(owner: string, id: string, options: { force?: boolean } = {}) {
  if (!options.force && activeJobs.has(id)) {
    throw createError({ statusCode: 409, statusMessage: 'Stop this batch before deleting it' })
  }
  markShotQueueCancelled(id)
  return mutate(owner, (queues) => {
    const index = queues.findIndex(item => item.id === id)
    if (index < 0) throw createError({ statusCode: 404, statusMessage: 'Queue not found' })
    const [removed] = queues.splice(index, 1)
    return removed
  })
}

export async function clearShotQueues(owner: string) {
  return mutate(owner, (queues) => {
    const kept: ShotQueue[] = []
    let removed = 0
    for (const queue of queues) {
      if (activeJobs.has(queue.id)) {
        kept.push(queue)
        continue
      }
      removed += 1
    }
    queues.splice(0, queues.length, ...kept)
    return { removed, remaining: kept.length }
  })
}

export function listOwnersWithQueues() {
  const root = join(libraryRoot(), 'users')
  if (!existsSync(root)) return []
  return readdirSync(root, { withFileTypes: true })
    .filter(entry => entry.isDirectory())
    .map(entry => entry.name)
    .filter(owner => existsSync(queuesPath(owner)))
}

export function remainingFromQueue(queue: ShotQueue): QueuedExtension[] {
  const family = persistPromptWrappers(queue)
  return queue.segments
    .filter(segment => segment.index > 0)
    .map(segment => ({
      prompt: segment.prompt,
      duration: segment.duration,
      permanenceRefs: segment.permanenceRefs,
      ...persistLoraFields(segment.loraStack || segment.loraName),
      ...family,
      ...persistPromptWrappers(segment)
    }))
}

export function lastCompletedIndex(queue: ShotQueue) {
  const done = queue.segments.filter(segment => segment.status === 'complete')
  if (!done.length) return -1
  return Math.max(...done.map(segment => segment.index))
}

export function liveSegmentPrompt(queue: ShotQueue, extensionIndex: number) {
  const segment = queue.segments.find(item => item.index === extensionIndex + 1)
  return {
    prompt: segment?.prompt || '',
    duration: segment?.duration || 5,
    permanenceRefs: segment?.permanenceRefs,
    ...persistLoraFields(segment?.loraStack || segment?.loraName),
    ...persistPromptWrappers(queue),
    ...persistPromptWrappers(segment)
  }
}