aigen/server/utils/promptComfy.ts

import { assertSharedGpu, sharedGpuHeaders, sharedGpuConfigured } from '~/server/utils/sharedGpu'
import { comfyInputFilename, getComfyHost } from '~/server/utils/comfy'
import { getSidecarImageHost } from '~/server/utils/imageComfy'
import { ensureSidecarReady } from '~/server/utils/imageComfyLifecycle'
import { buildVisionPromptWorkflow } from '~/server/utils/promptWorkflow'
import { normalizeShotScript } from '~/utils/parseRecommend'

export type PromptJobStatus = 'queued' | 'running' | 'complete' | 'error'

export interface PromptJob {
  id: string
  status: PromptJobStatus
  message: string
  text: string
  error?: string
  startedAt: number
}

const jobs = new Map<string, PromptJob>()
const MAX_JOBS = 20

export function getPromptComfyHost() {
  if (sharedGpuConfigured()) return getComfyHost()
  return getSidecarImageHost()
}

export function promptComfyConfigured() {
  return Boolean(getPromptComfyHost())
}

async function promptComfyFetch(path: string, init?: RequestInit) {
  const host = getPromptComfyHost()
  if (!host) {
    throw createError({ statusCode: 503, statusMessage: 'Qwen VL is not configured. Set COMFY_HOST or IMAGE_COMFY_HOST to Beast.' })
  }
  if (init?.method && !['GET', 'HEAD'].includes(init.method.toUpperCase())) {
    await assertSharedGpu()
    const headers = new Headers(init.headers)
    for (const [key, value] of Object.entries(sharedGpuHeaders())) headers.set(key, value)
    init = { ...init, headers }
  }
  try {
    return await fetch(`${host}${path}`, init)
  } catch (error) {
    throw createError({
      statusCode: 502,
      statusMessage: `Beast Qwen VL is unreachable (${host})`,
      data: { cause: error instanceof Error ? error.message : String(error) }
    })
  }
}

async function uploadPromptImage(file: { filename: string; data: Buffer; type?: string }, jobId?: string) {
  const body = new FormData()
  const blob = new Blob([new Uint8Array(file.data)], { type: file.type || 'application/octet-stream' })
  const filename = comfyInputFilename(file.filename, file.data, jobId)
  body.append('image', blob, filename)
  body.append('overwrite', 'true')
  body.append('type', 'input')
  const res = await promptComfyFetch('/upload/image', { method: 'POST', body })
  if (!res.ok) {
    throw createError({ statusCode: 502, statusMessage: `Beast image upload failed (${res.status})` })
  }
  const uploaded = (await res.json()) as { name: string; subfolder?: string; type?: string }
  return uploaded.name || filename
}

async function queuePromptGraph(graph: unknown, clientId: string) {
  const res = await promptComfyFetch('/prompt', {
    method: 'POST',
    headers: { 'Content-Type': 'application/json' },
    body: JSON.stringify({ prompt: graph, client_id: clientId })
  })
  const payload = await res.json().catch(() => ({}))
  if (!res.ok) {
    const raw = JSON.stringify(payload)
    if (/does not exist|class_type|not found/i.test(raw)) {
      throw createError({
        statusCode: 503,
        statusMessage: 'Beast ComfyUI rejected the Qwen VL graph. The GGUF VLM nodes are missing on this instance — Recommend will not use the sidecar.',
        data: payload
      })
    }
    const message = (payload as { error?: { message?: string } }).error?.message
      || (payload as { node_errors?: unknown }).node_errors
      || `Beast queue failed (${res.status})`
    throw createError({ statusCode: 502, statusMessage: String(message), data: payload })
  }
  return payload as { prompt_id: string }
}

const QWEN_VL_NODE_TYPES = ['GGUF_VLM_ModelLoader', 'GGUF_VLM_ImageToText', 'AigenPreviewText']

async function assertQwenVlWorkflow() {
  const res = await promptComfyFetch('/object_info')
  if (!res.ok) {
    throw createError({
      statusCode: 502,
      statusMessage: `Beast ComfyUI did not return node info (${res.status})`
    })
  }
  const info = await res.json().catch(() => null)
  if (!info || typeof info !== 'object') {
    throw createError({
      statusCode: 502,
      statusMessage: 'Beast ComfyUI returned unreadable node info'
    })
  }
  const missing = QWEN_VL_NODE_TYPES.filter(name => !(name in (info as Record<string, unknown>)))
  if (missing.length) {
    throw createError({
      statusCode: 503,
      statusMessage: `Beast ComfyUI is missing Qwen VL nodes (${missing.join(', ')}). Recommend is wired to Beast only and will not fall back to the sidecar.`
    })
  }
}

async function fetchPromptHistory(promptId: string) {
  const res = await promptComfyFetch(`/history/${encodeURIComponent(promptId)}`)
  if (!res.ok) return null
  return (await res.json()) as Record<string, unknown>
}

function collectStrings(value: unknown, depth = 0): string[] {
  if (value == null || depth > 8) return []
  if (typeof value === 'string') {
    const text = value.trim()
    return text ? [text] : []
  }
  if (typeof value === 'number' || typeof value === 'boolean') return []
  if (Array.isArray(value)) return value.flatMap(item => collectStrings(item, depth + 1))
  if (typeof value === 'object') {
    const record = value as Record<string, unknown>
    for (const key of ['text', 'string', 'strings']) {
      if (record[key] !== undefined) {
        const found = collectStrings(record[key], depth + 1)
        if (found.length) return found
      }
    }
    return Object.values(record).flatMap(item => collectStrings(item, depth + 1))
  }
  return []
}

function extractPromptText(history: Record<string, unknown> | null, promptId: string) {
  if (!history) return ''
  const wrapped = history[promptId] as { outputs?: Record<string, unknown> } | undefined
  const outputs = wrapped?.outputs || (history as { outputs?: Record<string, unknown> }).outputs
  if (!outputs || typeof outputs !== 'object') return ''
  const node4 = outputs['4']
  const candidates = [
    ...(node4 ? collectStrings(node4) : []),
    ...collectStrings(outputs)
  ].sort((a, b) => b.length - a.length)
  return candidates[0] || ''
}

function historyError(history: Record<string, unknown> | null, promptId: string) {
  const entry = history?.[promptId] as {
    status?: {
      status_str?: string
      completed?: boolean
      messages?: Array<[string, Record<string, unknown>]>
    }
  } | undefined
  const status = entry?.status?.status_str
  const err = entry?.status?.messages?.find(([type]) => type === 'execution_error')?.[1]
  const message = String(err?.exception_message || err?.message || '').trim()
  if (status === 'interrupted') return 'Prompt recommend was interrupted.'
  if (status === 'error') return message || 'Beast Qwen VL reported an execution error'
  if (entry?.status?.completed && !extractPromptText(history, promptId)) {
    return 'Beast Qwen VL finished without prompt text'
  }
  return null
}

async function waitForPromptText(promptId: string, timeoutMs = 240_000) {
  const started = Date.now()
  while (Date.now() - started < timeoutMs) {
    const history = await fetchPromptHistory(promptId)
    const text = extractPromptText(history, promptId)
    if (text) return text
    const error = historyError(history, promptId)
    if (error) {
      throw createError({ statusCode: 502, statusMessage: error })
    }
    await new Promise(resolve => setTimeout(resolve, 1200))
  }
  throw createError({ statusCode: 504, statusMessage: 'Beast Qwen VL timed out while writing the prompt' })
}

export function createPromptJob(): PromptJob {
  const job: PromptJob = {
    id: crypto.randomUUID(),
    status: 'queued',
    message: 'Queued',
    text: '',
    startedAt: Date.now()
  }
  jobs.set(job.id, job)
  while (jobs.size > MAX_JOBS) {
    const oldest = jobs.keys().next().value
    if (oldest) jobs.delete(oldest)
  }
  return job
}

export function getPromptJob(id: string) {
  return jobs.get(id)
}

export function promptJobSnapshot(job: PromptJob) {
  return {
    jobId: job.id,
    status: job.status,
    message: job.message,
    text: job.text,
    error: job.error,
    elapsedMs: Date.now() - job.startedAt
  }
}

function draftUserPrompt(draft: string, mode: string) {
  const studio = mode === 'edit'
    ? 'image-edit'
    : mode === 'shot-script'
      ? 'video-shot-script'
      : 'video'
  const format = studio === 'image-edit'
    ? 'If this is a single edit, output only the refined prompt. If it needs phases, use ---pass 1--- blocks.'
    : 'Output a MiniMax shot list. Each shot starts with "shot N" on its own line, then [SHOT CONFIGURATION], [SUBJECT DIRECTION & ACTION], and [AUDIO CUES], each on its own line. Blank line between shots. No markdown fences. No one-paragraph dump.'
  return `Studio mode: ${studio}\n${format}\n\nDraft Prompt:\n${draft}`
}

export async function runPromptRecommend(job: PromptJob, params: {
  image: { filename: string; data: Buffer; type?: string }
  draftPrompt: string
  systemPrompt: string
  mode: string
}) {
  job.status = 'running'
  job.message = 'Checking Beast Qwen VL…'
  await ensureSidecarReady((status) => {
    job.message = status.message
  })
  job.message = 'Confirming Qwen VL nodes on Beast…'
  await assertQwenVlWorkflow()
  job.message = 'Uploading still to Beast Qwen VL…'
  const imageName = await uploadPromptImage(params.image, job.id)
  job.message = 'Qwen VL is reading the still…'
  const graph = buildVisionPromptWorkflow({
    imageName,
    draftPrompt: draftUserPrompt(params.draftPrompt, params.mode),
    systemPrompt: params.systemPrompt
  })
  const queued = await queuePromptGraph(graph, job.id)
  job.message = 'Writing the recommended prompt…'
  const text = await waitForPromptText(queued.prompt_id)
  const trimmed = text.trim()
  job.text = params.mode === 'edit' ? trimmed : normalizeShotScript(trimmed)
  job.status = 'complete'
  job.message = 'Recommended prompt ready'
}

export function failPromptJob(job: PromptJob, error: unknown) {
  const err = error as { statusMessage?: string; message?: string }
  job.status = 'error'
  job.error = err.statusMessage || err.message || 'Prompt recommend failed'
  job.message = job.error
}

export function hasActivePromptJobs() {
  return [...jobs.values()].some(job => job.status === 'queued' || job.status === 'running')
}