import { normalizeImageIterations, resolveImageIteration } from '~/utils/imageIterations'
import { createJob, emitJob, type Job } from '~/server/utils/jobs'
import { ensureComfyReady } from '~/server/utils/comfyLifecycle'
import { assertImageScaleToTotalPixelsNode, comfyImagePrefix, getComfyHost, uploadImage, queuePrompt } from '~/server/utils/comfy'
import { withImageComfyHost, waitForImageEdit, downloadEditedImage, purgeImageComfyArtifacts } from '~/server/utils/imageComfy'
import { buildImageV2Workflow, IMAGE_V2_NODE_LABELS, resolveKreaGenerateAssets } from '~/server/utils/imageWorkflowV2'
import { ensureComfyLoraNames, persistLoraFields } from '~/server/utils/loras'
import { imageDimensions } from '~/server/utils/resolution'
import { emitChainJob } from '~/server/utils/watch'
import { saveStill, stillChainName } from '~/server/utils/library'
import { comfyJobPrefix } from '~/utils/outputNames'
import { restorePromptParts, wrappedPromptForComfy } from '~/utils/promptParts'
import { isXaigenStudio } from '~/utils/videoModels'
import {
  imageV2GraphMode,
  imageV2StackSpecials,
  mergeLegacyImageV2Loras,
  parseImageV2Engine,
  type ImageV2Engine,
  type ImageV2Mode,
  type ImageV2Task
} from '~/utils/imageV2'
import type { EditImageFile } from '~/server/utils/imageChain'

export type EditV2RunParams = {
  mode: ImageV2Mode
  engine?: ImageV2Engine
  task: ImageV2Task
  image?: EditImageFile | null
  reference: EditImageFile | null
  mask?: EditImageFile | null
  prompt: string
  negative: string
  steps: number
  seed: number
  cfg: number
  snofsModel: number
  snofsClip: number
  consistencyModel: number
  consistencyClip: number
  megapixels: number
  strength?: number
  width?: number
  height?: number
  turbo?: boolean
  aspect?: string
  sourceStillId?: string
  referenceStillId?: string
  loraStack?: import('~/utils/loras').LoraStackItem[]
  passes?: import('~/utils/imageIterations').ImageIteration[]
  passMode?: 'batch' | 'chain'
}

export async function runEditV2(job: Job, params: EditV2RunParams) {
  const library = job.library
  if (!library) throw new Error('Edit job is missing library metadata')
  const graphMode = imageV2GraphMode(params.mode, {
    hasImage: Boolean(params.image),
    hasReference: Boolean(params.reference)
  })
  if (graphMode === 'refine' && !params.mask) {
    throw new Error('Refine requires a mask. Refusing to fall back to Edit.')
  }
  if (graphMode === 'compose' && !params.reference) {
    throw new Error('Compose requires image B. Refusing to fall back to one-image generation.')
  }
  if (graphMode === 'edit' && params.reference && params.mode !== 'iterate') {
    throw new Error('Edit mode takes one image. Use Compose for two stills.')
  }
  if (graphMode !== 'generate' && !params.image) {
    throw new Error('This v2 mode requires still A.')
  }

  const canChain = params.mode === 'iterate' || params.mode === 'edit' || params.mode === 'compose'
  const extraPasses = canChain
    ? normalizeImageIterations(params.passes)
    : []
  const prompts = [
    library.promptMid || restorePromptParts(params.prompt, {
      pre: library.promptPre,
      prompt: params.prompt,
      post: library.promptPost
    }).prompt,
    ...extraPasses.map(item => item.prompt || library.promptMid || library.prompt || params.prompt)
  ]
  const chainTotal = prompts.length
  // Text-only iterate is always independent generations. With a still, honor batch vs chain.
  const batch = graphMode === 'generate' || params.passMode !== 'chain'
  const unit = params.mode === 'iterate' ? 'Iteration' : (batch ? 'Batch' : 'Pass')
  library.chainTotal = chainTotal
  library.familyId = chainTotal > 1 ? (library.familyId || crypto.randomUUID()) : library.familyId

  try {
    // Keep the studio slot across all iterations, including upload/save gaps.
    library.chainContinuing = true
    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'
      })
    }, { skipBusyWait: true })
    if (job.status === 'cancelled') throw new Error('Job interrupted.')
    job.imageComfyHost = getComfyHost()

    await withImageComfyHost(job.imageComfyHost, async () => {
      const generate = graphMode === 'generate'
      const engine = parseImageV2Engine(params.engine, 'flux')
      const engineLabel = engine === 'krea' ? 'Krea' : 'Flux'
      const { applyImageFill } = await import('~/server/utils/imagePad')
      const { composeFillPrompt } = await import('~/utils/imageAspect')
      let current = params.image || null
      let startPadded = false
      let mask = params.mask
      let parentStillId: string | undefined
      let uploadedRef: { name: string } | null = null
      let uploadedMask: { name: string } | null = null

      await ensureComfyLoraNames('image')
      if (!generate && params.megapixels > 0) await assertImageScaleToTotalPixelsNode()
      const krea = engine === 'krea' ? await resolveKreaGenerateAssets() : null
      if (!generate && graphMode === 'compose' && params.reference) {
        uploadedRef = await uploadImage({
          ...params.reference,
          filename: `ref_${params.reference.filename || 'image_b.png'}`
        }, job.id)
      }

      for (let index = 0; index < prompts.length; index++) {
        if (job.status === 'cancelled') throw new Error('Job interrupted.')
        if (index > 0 && library.stopAfterCurrent === true) break
        const last = index === prompts.length - 1
        const pass = index > 0 ? extraPasses[index - 1] : undefined
        const { steps, cfg, seed, strength } = resolveImageIteration(params, pass)
        const turbo = pass?.steps !== undefined || pass?.cfg !== undefined ? false : params.turbo === true
        const appliedStack = mergeLegacyImageV2Loras(pass?.loraStack ?? params.loraStack ?? library.loraStack, {
          engine,
          mode: graphMode,
          xaigen: isXaigenStudio(),
          snofsModel: pass?.loraStack !== undefined ? 0 : params.snofsModel,
          snofsClip: pass?.loraStack !== undefined ? 0 : params.snofsClip,
          consistencyModel: pass?.loraStack !== undefined ? 0 : params.consistencyModel,
          consistencyClip: pass?.loraStack !== undefined ? 0 : params.consistencyClip,
          kreaConcept: krea?.conceptLora
        })
        const appliedSpecials = imageV2StackSpecials(appliedStack, engine)
        const snofsModel = appliedSpecials.snofs?.strengthModel ?? (pass?.loraStack !== undefined ? 0 : params.snofsModel)
        const snofsClip = appliedSpecials.snofs?.strengthClip ?? (pass?.loraStack !== undefined ? 0 : params.snofsClip)
        const consistencyModel = appliedSpecials.consistency?.strengthModel ?? (pass?.loraStack !== undefined ? 0 : params.consistencyModel)
        const consistencyClip = appliedSpecials.consistency?.strengthClip ?? (pass?.loraStack !== undefined ? 0 : params.consistencyClip)

        library.steps = steps
        library.cfg = cfg
        job.maxStep = steps
        library.prompt = prompts[index]
        library.seed = seed
        library.chainIndex = index
        library.chainStep = index + 1
        library.chainLabel = chainTotal > 1 ? `${unit} ${index + 1}` : undefined
        const passName = stillChainName(library.name || '', index)
        let prompt = prompts[index]
        let image = current
        let padded = false

        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'
            })
          }, { skipBusyWait: true })
        }

        emitChainJob(job, {
          type: 'status',
          message: generate
            ? `Queueing ${engineLabel} iterate generate on Beast...`
            : graphMode === 'refine'
              ? 'Uploading canvas and mask to Beast...'
              : graphMode === 'compose' ? 'Uploading stills A and B to Beast...' : 'Uploading still A to Beast...',
          progress: generate ? 12 : 8
        })

        if (index === 0 && !generate && image) {
          const filled = await applyImageFill({
            image,
            mask: graphMode === 'refine' ? mask : null,
            aspect: params.aspect || 'auto'
          })
          image = filled.image
          current = filled.image
          if (filled.mask) {
            mask = filled.mask
            uploadedMask = await uploadImage({
              ...filled.mask,
              filename: `mask_${filled.mask.filename || 'refine-mask.png'}`
            }, job.id)
          }
          padded = filled.padded
          startPadded = filled.padded
          prompt = composeFillPrompt(prompt, padded)
          if (padded) {
            emitChainJob(job, {
              type: 'status',
              message: `Filling to ${filled.aspect} without stretching...`,
              progress: 7
            })
          }
        } else if (batch && startPadded) {
          prompt = composeFillPrompt(prompt, true)
        }

        if (index === 0 && graphMode === 'refine' && mask && !uploadedMask) {
          uploadedMask = await uploadImage({
            ...mask,
            filename: `mask_${mask.filename || 'refine-mask.png'}`
          }, job.id)
        }

        const uploaded = generate || !image ? null : await uploadImage(image, job.id)
        if (job.status === 'cancelled') throw new Error('Job interrupted.')

        if (!generate) {
          emitChainJob(job, {
            type: 'status',
            message: params.mode === 'iterate'
              ? `Queueing ${engineLabel} iterate (${graphMode}) on Beast...`
              : graphMode === 'refine'
                ? `Queueing ${engineLabel} v2 refine on Beast...`
                : graphMode === 'compose'
                  ? `Queueing ${engineLabel} v2 compose (${params.task}) on Beast...`
                  : `Queueing ${engineLabel} v2 edit on Beast...`,
            progress: 12
          })
        }

        const built = buildImageV2Workflow({
          mode: graphMode,
          engine,
          task: generate ? 't2i' : params.task,
          prompt: wrappedPromptForComfy(library, prompt, index === 0 ? params.prompt : prompt),
          negative: params.negative,
          imageAName: uploaded?.name,
          imageBName: uploadedRef?.name,
          maskName: uploadedMask?.name,
          strength,
          snofsModel,
          snofsClip,
          consistencyModel,
          consistencyClip,
          steps,
          cfg,
          seed,
          megapixels: params.megapixels,
          width: params.width,
          height: params.height,
          aspect: params.aspect,
          turbo,
          filenamePrefix: chainTotal > 1 ? `${comfyImagePrefix()}_${comfyJobPrefix(job.id)}_p${index + 1}` : `${comfyImagePrefix()}_${comfyJobPrefix(job.id)}`,
          kreaUnetName: krea?.unet,
          kreaClipName: krea?.clip,
          kreaVaeName: krea?.vae,
          kreaConceptLora: krea?.conceptLora,
          loraStack: appliedStack
        })
        const queued = await queuePrompt(built.graph, job.clientId)
        job.promptId = queued.prompt_id
        job.status = 'running'
        emitChainJob(job, {
          type: 'status',
          message: `Running ${built.workflowFile}...`,
          progress: 18,
          maxStep: steps
        })

        const output = await waitForImageEdit({
          promptId: queued.prompt_id,
          clientId: job.clientId,
          timeoutMs: 10 * 60 * 1000,
          engineLabel,
          nodeLabel: (node) => {
            if (node === '15' || node === '19') return `Sampling ${engineLabel}`
            return IMAGE_V2_NODE_LABELS[node]
          },
          onProgress: (event) => {
            emitChainJob(job, {
              type: 'status',
              message: event.message,
              progress: event.progress,
              step: event.step,
              maxStep: event.maxStep || steps,
              node: event.node
            })
          },
          isCancelled: () => job.status === 'cancelled'
        })

        emitChainJob(job, { type: 'status', message: `Saving ${engineLabel} v2 still...`, progress: 94 })
        job.saving = true
        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,
          promptPre: library.promptPre,
          promptPost: library.promptPost,
          familyId: library.familyId,
          parentStillId: batch ? (params.sourceStillId || library.stillId) : parentStillId,
          chainIndex: index,
          settings: {
            kind: 'imagev2',
            mode: params.mode,
            engine,
            task: generate ? 't2i' : params.task,
            negative: params.negative,
            snofsModel,
            snofsClip,
            consistencyModel,
            consistencyClip,
            steps,
            cfg,
            megapixels: generate || !(params.megapixels > 0) ? undefined : params.megapixels,
            scaleToTotalPixels: generate ? undefined : params.megapixels > 0,
            turbo,
            strength: generate ? undefined : strength,
            width: generate ? params.width : undefined,
            height: generate ? params.height : undefined,
            aspect: params.aspect || 'auto',
            seed,
            sourceStillId: params.sourceStillId || library.stillId,
            referenceStillId: params.referenceStillId,
            ...persistLoraFields(appliedStack)
          }
        })
        job.stillId = still?.id
        if (!batch) parentStillId = still?.id
        await purgeImageComfyArtifacts({
          output: { filename: output.filename, subfolder: output.subfolder, type: output.type },
          inputName: uploaded?.name,
          inputSubfolder: uploaded?.subfolder,
          extraInputNames: last
            ? [uploadedRef?.name, uploadedMask?.name].filter((name): name is string => Boolean(name))
            : [],
          promptId: job.promptId
        })

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

        job.status = 'complete'
        emitChainJob(job, {
          type: 'complete',
          message: chainTotal > 1
            ? `${engineLabel} ${params.mode === 'iterate' ? 'iterate' : `v2 ${batch ? 'batch' : 'chain'}`} finished on Beast`
            : `${engineLabel} v2 finished on Beast`,
          progress: 100,
          stillId: still?.id,
          filename: output.filename,
          subfolder: output.subfolder,
          mediaType: 'image',
          hideThumbnail: job.hideThumbnail,
          folderLocked: library.folderLocked
        })
      }
    })
  } catch (error) {
    if (job.status !== 'error' && job.status !== 'cancelled') {
      const message = error instanceof Error ? error.message : String(error)
      job.status = 'error'
      job.error = message
      emitJob(job, { type: 'error', error: message, message })
    }
  } finally {
    job.saving = false
    library.chainContinuing = false
    const { onLiveVideoSettled } = await import('~/server/utils/studioQueue')
    await onLiveVideoSettled(job)
  }
}