import { ensureComfyReady } from '~/server/utils/comfyLifecycle'
import { getComfyHost, comfyConfigured, uploadImage, queuePrompt, purgeComfyArtifacts } from '~/server/utils/comfy'
import { withImageComfyHost, waitForImageEdit, downloadEditedImage } from '~/server/utils/imageComfy'
import { buildEditWorkflow } from '~/server/utils/imageWorkflow'
import { imageDimensions } from '~/server/utils/resolution'

function parseSteps(raw: unknown) {
  const value = Number(raw)
  if (!Number.isFinite(value)) return 20
  return Math.min(40, Math.max(1, Math.round(value)))
}

function parseCfg(raw: unknown) {
  const value = Number(raw)
  if (!Number.isFinite(value)) return 1
  return Math.min(8, Math.max(0.5, Math.round(value * 10) / 10))
}

function parsePasses(raw: string | undefined) {
  if (!raw) return [] as { prompt: string }[]
  try {
    const parsed = JSON.parse(raw)
    if (!Array.isArray(parsed)) return []
    return parsed
      .map((item: { prompt?: unknown }) => ({ prompt: String(item?.prompt || '').trim() }))
      .filter(item => item.prompt)
  } catch {
    return []
  }
}

export default defineEventHandler(async (event) => {
  const form = await readMultipartFormData(event).catch(() => null)
  const fields: Record<string, string> = {}
  let image: { filename: string; data: Buffer; type?: string } | null = null
  let reference: { filename: string; data: Buffer; type?: string } | null = null
  for (const part of form || []) {
    if (part.name === 'image' && part.filename && part.data?.length) {
      image = { filename: part.filename, data: part.data, type: part.type }
    } else if ((part.name === 'image2' || part.name === 'refImage') && part.filename && part.data?.length) {
      reference = { filename: part.filename, data: part.data, type: part.type }
    } else if (part.name && part.data) {
      fields[part.name] = part.data.toString('utf8')
    }
  }

  const prompt = (fields.prompt || '').trim()
  const extraPasses = parsePasses(fields.passes)
  if (!image) {
    throw createError({ statusCode: 400, statusMessage: 'An input image is required' })
  }
  if (!prompt) {
    throw createError({ statusCode: 400, statusMessage: 'An edit prompt is required' })
  }
  if (extraPasses.some(item => !item.prompt)) {
    throw createError({ statusCode: 400, statusMessage: 'Each extra pass needs a prompt' })
  }
  if (!comfyConfigured()) {
    throw createError({
      statusCode: 503,
      statusMessage: 'Beast ComfyUI is not configured. Set COMFY_HOST.'
    })
  }

  const ownerKey = libraryOwnerKey(event)
  const library = publicLibrary(event)
  const folderId = library.folders.some(folder => folder.id === fields.folderId)
    ? fields.folderId
    : library.folders[0]?.id
  if (!folderId) {
    throw createError({ statusCode: 400, statusMessage: 'Create a library folder before editing' })
  }
  assertFolderExists(event, folderId)

  const destFolder = library.folders.find(folder => folder.id === folderId)
  const folderLocked = Boolean(destFolder?.protected && !destFolder.unlocked)
  const hideInput = fields.hideInput === 'true'
  const hideThumbnail = fields.hideThumbnail === 'true'
  const steps = parseSteps(fields.steps)
  const cfg = parseCfg(fields.cfg)
  const seed = fields.seed && fields.seed !== 'random'
    ? Number(fields.seed)
    : Math.floor(Math.random() * 2_147_483_647)
  const size = imageDimensions(image.data)

  const still = await rememberInputStill({
    ownerKey,
    folderId,
    filename: image.filename,
    data: image.data,
    width: size?.width || 0,
    height: size?.height || 0,
    hideInput
  })
  if (reference) {
    await rememberInputStill({
      ownerKey,
      folderId,
      filename: reference.filename,
      data: reference.data,
      hideInput
    })
  }

  const passes = extraPasses
  const chainTotal = 1 + passes.length
  const familyId = chainTotal > 1 ? crypto.randomUUID() : undefined
  const job = createJob()
  job.kind = 'edit'
  job.maxStep = steps
  job.hideThumbnail = hideThumbnail
  job.library = {
    ownerKey,
    folderId,
    hideThumbnail,
    hideInput,
    folderLocked,
    name: (fields.name || '').trim().slice(0, 80),
    prompt,
    aspect: 'auto',
    width: size?.width || 0,
    height: size?.height || 0,
    steps,
    turbo: true,
    seed,
    cfg,
    stillId: still?.id,
    familyId,
    chainIndex: 0,
    chainStep: 1,
    chainTotal,
    chainLabel: chainTotal > 1 ? 'Pass 1' : undefined,
    passes
  }
  emitJob(job, {
    type: 'status',
    message: 'Checking Beast ComfyUI...',
    progress: 2,
    chainStep: 1,
    chainTotal,
    chainLabel: job.library.chainLabel
  })

  void runEdit(job, {
    image,
    reference,
    prompt,
    passes,
    negative: (fields.negative || '').trim(),
    steps,
    seed,
    cfg
  }).catch(async (error) => {
    if (job.status === 'error' || job.status === 'cancelled') return
    const message = error instanceof Error ? error.message : String(error)
    job.status = 'error'
    job.error = message
    emitJob(job, { type: 'error', error: message, message })
  })

  return {
    jobId: job.id,
    seed,
    steps,
    hideThumbnail,
    folderLocked,
    dual: Boolean(reference),
    chainTotal
  }
})

async function runEdit(
  job: ReturnType<typeof createJob>,
  params: {
    image: { filename: string; data: Buffer; type?: string }
    reference: { filename: string; data: Buffer; type?: string } | null
    prompt: string
    passes: { prompt: string }[]
    negative: string
    steps: number
    seed: number
    cfg: number
  }
) {
  const library = job.library
  if (!library) throw new Error('Edit job is missing library metadata')
  const prompts = [params.prompt, ...params.passes.map(item => item.prompt)]
  const chainTotal = prompts.length
  library.chainTotal = chainTotal
  library.familyId = chainTotal > 1 ? (library.familyId || crypto.randomUUID()) : library.familyId

  await ensureComfyReady((status) => {
    emitChainJob(job, {
      type: status.state === 'busy' ? 'busy' : 'status',
      message: status.message,
      progress: status.state === 'online' ? Math.max(job.progress, 6) : Math.max(job.progress, 3),
      busy: status.state === 'busy'
    })
  })
  if (job.status === 'cancelled') throw new Error('Job interrupted.')
  job.imageComfyHost = getComfyHost()

  await withImageComfyHost(job.imageComfyHost, async () => {
    let current = params.image
    let parentStillId: string | undefined
    for (let index = 0; index < prompts.length; index++) {
      if (job.status === 'cancelled') throw new Error('Job interrupted.')
      const prompt = prompts[index]
      const last = index === prompts.length - 1
      if (index > 0) {
        await ensureComfyReady((status) => {
          emitChainJob(job, {
            type: status.state === 'busy' ? 'busy' : 'status',
            message: status.message,
            progress: status.state === 'online' ? 4 : 2,
            busy: status.state === 'busy'
          })
        })
      }
      const seed = index === 0 ? params.seed : Math.floor(Math.random() * 2_147_483_647)
      library.prompt = prompt
      library.seed = seed
      library.chainIndex = index
      library.chainStep = index + 1
      library.chainLabel = chainTotal > 1 ? `Pass ${index + 1}` : undefined
      const reference = index === 0 ? params.reference : null
      const dual = Boolean(reference)
      const passName = stillChainName(library.name || '', index)

      emitChainJob(job, {
        type: 'status',
        message: dual ? 'Uploading both stills to Beast...' : 'Uploading still to Beast...',
        progress: 8
      })
      const uploaded = await uploadImage(current, job.id)
      const uploadedRef = reference
        ? await uploadImage({
          ...reference,
          filename: `ref_${reference.filename || 'image2.png'}`
        }, job.id)
        : null
      if (job.status === 'cancelled') throw new Error('Job interrupted.')

      emitChainJob(job, {
        type: 'status',
        message: dual
          ? 'Queueing Flux.2 Klein two-image edit on Beast...'
          : 'Queueing Flux.2 Klein edit on Beast...',
        progress: 12
      })
      const graph = buildEditWorkflow({
        prompt,
        negative: params.negative,
        imageName: uploaded.name,
        referenceImageName: uploadedRef?.name,
        steps: params.steps,
        seed,
        cfg: params.cfg,
        filenamePrefix: `aigen_edit_${job.id.slice(0, 8)}_p${index + 1}`
      })
      const queued = await queuePrompt(graph, job.clientId)
      job.promptId = queued.prompt_id
      job.status = 'running'
      emitChainJob(job, {
        type: 'status',
        message: dual ? 'Applying image 2 onto image 1...' : 'Editing still on Beast...',
        progress: 18,
        maxStep: params.steps
      })

      const output = await waitForImageEdit({
        promptId: queued.prompt_id,
        clientId: job.clientId,
        timeoutMs: 10 * 60 * 1000,
        onProgress: (event) => {
          emitChainJob(job, {
            type: 'status',
            message: event.message,
            progress: event.progress,
            step: event.step,
            maxStep: event.maxStep || params.steps,
            node: event.node
          })
        },
        isCancelled: () => job.status === 'cancelled'
      })

      emitChainJob(job, { type: 'status', message: 'Saving edited still...', progress: 94 })
      const buffer = await downloadEditedImage(output)
      const size = imageDimensions(buffer)
      const still = await saveStill({
        ownerKey: library.ownerKey,
        folderId: library.folderId,
        filename: passName ? `${passName}.png` : output.filename,
        data: buffer,
        width: size?.width || 0,
        height: size?.height || 0,
        hideInput: job.hideThumbnail === true,
        role: 'output',
        name: passName || undefined,
        prompt,
        familyId: library.familyId,
        parentStillId,
        chainIndex: index
      })
      job.stillId = still?.id
      parentStillId = still?.id
      await purgeComfyArtifacts({
        video: { filename: output.filename, subfolder: output.subfolder, type: output.type },
        imageName: uploaded.name,
        imageSubfolder: uploaded.subfolder,
        extraImageNames: uploadedRef?.name ? [uploadedRef.name] : [],
        promptId: job.promptId
      })

      if (!last) {
        emitChainJob(job, {
          type: 'checkpoint',
          message: `Pass ${index + 1} saved`,
          progress: 100,
          stillId: still?.id,
          hideThumbnail: job.hideThumbnail,
          folderLocked: library.folderLocked
        })
        current = {
          filename: still?.filename || `pass_${index + 1}.png`,
          data: buffer,
          type: 'image/png'
        }
        continue
      }

      job.status = 'complete'
      emitChainJob(job, {
        type: 'complete',
        message: chainTotal > 1 ? 'Edit chain finished on Beast' : 'Edit finished on Beast',
        progress: 100,
        stillId: still?.id,
        filename: output.filename,
        subfolder: output.subfolder,
        mediaType: 'image',
        hideThumbnail: job.hideThumbnail,
        folderLocked: library.folderLocked
      })
    }
  })
}