aigen/server/utils/imageComfy.ts

import { AsyncLocalStorage } from 'node:async_hooks'
import { getComfyHost } from '~/server/utils/comfy'

const imageHostAls = new AsyncLocalStorage<string>()
let imageComfyHostOverride = ''

export type ImageEditBox = 'beast' | 'sidecar'

export function withImageComfyHost<T>(host: string, fn: () => Promise<T>) {
  return imageHostAls.run(String(host || '').replace(/\/$/, ''), fn)
}

export function setImageComfyHostOverride(url: string) {
  imageComfyHostOverride = String(url || '').replace(/\/$/, '')
}

function normalizeHost(raw: string, port = '') {
  let host = String(raw || '').trim().replace(/\/$/, '')
  if (!host) return ''
  if (!/^https?:\/\//i.test(host)) host = `http://${host}`
  try {
    const url = new URL(host)
    if (port && !url.port) url.port = port
    return url.origin
  } catch {
    return port ? `${host}:${port}` : host
  }
}

export function getBeastImageHost() {
  const config = useRuntimeConfig()
  const explicit = String(config.imageComfyPrimaryHost || process.env.IMAGE_COMFY_PRIMARY_HOST || '').trim()
  if (explicit) return normalizeHost(explicit)
  try {
    return getComfyHost()
  } catch {
    return ''
  }
}

export function getSidecarImageHost() {
  const config = useRuntimeConfig()
  const fallback = String(config.imageComfyFallbackHost || process.env.IMAGE_COMFY_FALLBACK_HOST || '').trim()
  if (fallback) return normalizeHost(fallback)
  const legacy = String(config.imageComfyHost || process.env.IMAGE_COMFY_HOST || '').trim()
  const port = String(config.imageComfyPort || process.env.IMAGE_COMFY_PORT || '').trim()
  return normalizeHost(legacy, port)
}

export function getImageComfyHost() {
  const fromAls = imageHostAls.getStore()
  if (fromAls) return fromAls
  if (imageComfyHostOverride) return imageComfyHostOverride
  return getBeastImageHost() || getSidecarImageHost()
}

export function imageComfyConfigured() {
  return Boolean(getBeastImageHost() || getSidecarImageHost())
}

export function sameImageHost(a: string, b: string) {
  return Boolean(a && b && a.replace(/\/$/, '') === b.replace(/\/$/, ''))
}

export async function imageComfyFetch(path: string, init?: RequestInit) {
  const host = getImageComfyHost()
  if (!host) {
    throw createError({ statusCode: 500, statusMessage: 'No image ComfyUI host is configured' })
  }
  try {
    return await fetch(`${host}${path}`, init)
  } catch (error) {
    throw createError({
      statusCode: 502,
      statusMessage: `Image ComfyUI host unreachable (${host})`,
      data: { cause: error instanceof Error ? error.message : String(error) }
    })
  }
}

export function imageComfyWsUrl(clientId: string) {
  return `${getImageComfyHost().replace(/^http/, 'ws')}/ws?clientId=${encodeURIComponent(clientId)}`
}

function imageInputFilename(original: string, jobId?: string) {
  const raw = String(original || 'still.png')
  const dot = raw.lastIndexOf('.')
  const ext = (dot >= 0 ? raw.slice(dot) : '.png').replace(/[^.a-zA-Z0-9]/g, '') || '.png'
  const base = (dot >= 0 ? raw.slice(0, dot) : raw).replace(/[^a-zA-Z0-9._-]+/g, '_').slice(0, 48) || 'still'
  const id = String(jobId || crypto.randomUUID()).replace(/-/g, '').slice(0, 10)
  return `aigen_edit_${id}_${base}${ext}`
}

export async function uploadImageEdit(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 = imageInputFilename(file.filename, jobId)
  body.append('image', blob, filename)
  body.append('overwrite', 'true')
  body.append('type', 'input')
  const res = await imageComfyFetch('/upload/image', { method: 'POST', body })
  if (!res.ok) {
    throw createError({ statusCode: 502, statusMessage: `Image upload failed (${res.status})` })
  }
  const uploaded = (await res.json()) as { name: string; subfolder?: string; type?: string }
  return {
    name: uploaded.name || filename,
    subfolder: uploaded.subfolder || '',
    type: uploaded.type || 'input'
  }
}

export async function queueImagePrompt(graph: unknown, clientId: string) {
  const res = await imageComfyFetch('/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 message = (payload as { error?: { message?: string } }).error?.message
      || (payload as { node_errors?: unknown }).node_errors
      || `Queue failed (${res.status})`
    throw createError({ statusCode: 502, statusMessage: String(message), data: payload })
  }
  return payload as { prompt_id: string; number?: number }
}

export async function interruptImageComfy(host?: string) {
  if (!imageComfyConfigured() && !host) return false
  const run = async () => {
    const res = await imageComfyFetch('/interrupt', { method: 'POST' })
    return res.ok
  }
  if (host) return withImageComfyHost(host, run)
  return run()
}

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

function isStillFile(item: { filename?: string } | null | undefined) {
  if (!item?.filename) return false
  return /\.(png|jpe?g|webp)$/i.test(item.filename)
}

function findStill(value: unknown, depth = 0): { filename: string; subfolder: string; type: string } | null {
  if (!value || typeof value !== 'object' || depth > 8) return null
  if (Array.isArray(value)) {
    const files = value.filter((item): item is { filename?: string; subfolder?: string; type?: string } => Boolean(item && typeof item === 'object'))
    const match = files.find(isStillFile)
    if (match?.filename) {
      return {
        filename: String(match.filename),
        subfolder: String(match.subfolder || ''),
        type: String(match.type || 'output')
      }
    }
    for (const item of value) {
      const nested = findStill(item, depth + 1)
      if (nested) return nested
    }
    return null
  }
  const record = value as { filename?: string; subfolder?: string; type?: string; images?: unknown }
  if (record.filename && isStillFile(record)) {
    return {
      filename: String(record.filename),
      subfolder: String(record.subfolder || ''),
      type: String(record.type || 'output')
    }
  }
  for (const nested of Object.values(value as Record<string, unknown>)) {
    const found = findStill(nested, depth + 1)
    if (found) return found
  }
  return null
}

export function extractEditedImage(history: Record<string, unknown> | null, promptId: string) {
  if (!history) return null
  const wrapped = history[promptId] as { outputs?: Record<string, unknown> } | undefined
  if (wrapped) return findStill(wrapped.outputs || {}) || findStill(wrapped)
  if ((history as { outputs?: unknown }).outputs) {
    return findStill((history as { outputs?: unknown }).outputs) || findStill(history)
  }
  return findStill(history)
}

export async function downloadEditedImage(image: { filename: string; subfolder: string; type: string }) {
  const params = new URLSearchParams({
    filename: image.filename,
    subfolder: image.subfolder || '',
    type: image.type || 'output'
  })
  const res = await imageComfyFetch(`/view?${params.toString()}`)
  if (!res.ok) {
    throw createError({ statusCode: 502, statusMessage: `Failed to fetch edited still (${res.status})` })
  }
  return Buffer.from(await res.arrayBuffer())
}

function historyError(history: Record<string, unknown> | null, promptId: string) {
  const entry = history?.[promptId] as { status?: { status_str?: string; completed?: boolean } } | undefined
  const status = entry?.status?.status_str
  if (status === 'interrupted') return 'Job interrupted.'
  if (status === 'error') return 'Image ComfyUI reported an execution error'
  if (entry?.status?.completed && !extractEditedImage(history, promptId)) {
    return 'Image ComfyUI finished without an output still'
  }
  return null
}

export async function waitForImageEdit(opts: {
  promptId: string
  clientId: string
  timeoutMs?: number
  onProgress?: (event: { message: string; progress: number; step?: number; maxStep?: number; node?: string | null }) => void
  isCancelled?: () => boolean
}) {
  const timeoutMs = opts.timeoutMs || 180_000
  const started = Date.now()
  let settled = false
  let lastError: string | null = null
  let result: { filename: string; subfolder: string; type: string } | null = null

  const finish = (image: { filename: string; subfolder: string; type: string } | null, error?: string) => {
    if (settled) return
    settled = true
    lastError = error || null
    result = image
  }

  let ws: WebSocket | null = null
  try {
    ws = new WebSocket(imageComfyWsUrl(opts.clientId))
    ws.addEventListener('message', (event) => {
      const payload = JSON.parse(String(event.data || '{}')) as {
        type?: string
        data?: { value?: number; max?: number; node?: string | null; prompt_id?: string }
      }
      if (payload.type === 'progress' && payload.data) {
        const max = Math.max(1, Number(payload.data.max || 4))
        const value = Number(payload.data.value || 0)
        opts.onProgress?.({
          message: 'Sampling edit...',
          progress: Math.min(90, 20 + Math.round((value / max) * 60)),
          step: value,
          maxStep: max,
          node: payload.data.node || '8'
        })
      }
      if (payload.type === 'executing' && payload.data?.node) {
        opts.onProgress?.({
          message: payload.data.node === '8' ? 'Sampling edit...' : 'Running Qwen Edit...',
          progress: payload.data.node === '10' ? 92 : 30,
          node: payload.data.node
        })
      }
      if (payload.type === 'execution_error') {
        finish(null, 'Image ComfyUI reported an execution error')
      }
    })
  } catch {
    ws = null
  }

  while (!settled && Date.now() - started < timeoutMs) {
    if (opts.isCancelled?.()) {
      finish(null, 'Job interrupted.')
      break
    }
    const history = await fetchImageHistory(opts.promptId)
    const image = extractEditedImage(history, opts.promptId)
    if (image) {
      finish(image)
      break
    }
    const error = historyError(history, opts.promptId)
    if (error) {
      finish(null, error)
      break
    }
    await new Promise(resolve => setTimeout(resolve, 700))
  }

  try { ws?.close() } catch { /* ignore */ }
  if (result) return result
  throw createError({
    statusCode: 502,
    statusMessage: lastError || 'Image edit timed out waiting for Qwen Edit to finish'
  })
}

function purgeImageEnabled() {
  return useRuntimeConfig().purgeComfyOutputs !== false
}

async function deleteSidecarFile(file: { filename: string; subfolder?: string; type?: string }) {
  const payload = {
    filename: file.filename,
    subfolder: file.subfolder || '',
    type: file.type || 'output'
  }
  for (const path of ['/delete', '/aigen/purge']) {
    try {
      const res = await imageComfyFetch(path, {
        method: 'POST',
        headers: { 'Content-Type': 'application/json' },
        body: JSON.stringify(payload),
        signal: AbortSignal.timeout(4000)
      })
      if (res.ok) return true
    } catch {
      /* Comfy may not expose this route */
    }
  }
  return false
}

async function purgeImageOnDesktop(opts: {
  inputName?: string
  inputSubfolder?: string
  output?: { filename: string; subfolder: string; type: string }
}) {
  const config = useRuntimeConfig()
  const host = getImageComfyHost()
  const sidecar = getSidecarImageHost()
  const useSidecarAgent = sidecar && sameImageHost(host, sidecar)
  const controlUrl = String(
    useSidecarAgent
      ? (config.imageComfyControlUrl || process.env.IMAGE_COMFY_CONTROL_URL || '')
      : (config.comfyControlUrl || process.env.COMFY_CONTROL_URL || '')
  ).replace(/\/$/, '')
  const token = String(
    useSidecarAgent
      ? (config.imageComfyControlToken || process.env.IMAGE_COMFY_CONTROL_TOKEN || '')
      : (config.comfyControlToken || process.env.COMFY_CONTROL_TOKEN || '')
  )
  if (!controlUrl) return
  try {
    await fetch(`${controlUrl}/purge`, {
      method: 'POST',
      headers: {
        'Content-Type': 'application/json',
        Accept: 'application/json',
        ...(token ? { Authorization: `Bearer ${token}` } : {})
      },
      body: JSON.stringify({
        imageName: opts.inputName || '',
        imageSubfolder: opts.inputSubfolder || '',
        output: opts.output || null
      }),
      signal: AbortSignal.timeout(4000)
    })
  } catch {
    /* Sidecar host agent is optional; Comfy HTTP delete is the primary path. */
  }
}

export async function purgeImageComfyArtifacts(opts: {
  output?: { filename: string; subfolder: string; type: string }
  inputName?: string
  inputSubfolder?: string
  promptId?: string
}) {
  if (!purgeImageEnabled()) return
  const files: { filename: string; subfolder: string; type: string }[] = []
  if (opts.output?.filename) {
    files.push({
      filename: opts.output.filename,
      subfolder: opts.output.subfolder || '',
      type: opts.output.type || 'output'
    })
  }
  if (opts.inputName) {
    files.push({
      filename: opts.inputName,
      subfolder: opts.inputSubfolder || '',
      type: 'input'
    })
  }
  for (const file of files) {
    await deleteSidecarFile(file)
  }
  await purgeImageOnDesktop({
    inputName: opts.inputName,
    inputSubfolder: opts.inputSubfolder,
    output: opts.output
  })
  if (opts.promptId) {
    try {
      await imageComfyFetch('/history', {
        method: 'POST',
        headers: { 'Content-Type': 'application/json' },
        body: JSON.stringify({ delete: [opts.promptId] }),
        signal: AbortSignal.timeout(4000)
      })
    } catch {
      /* ignore */
    }
  }
}