// Nitro bundles these JSON graphs into the production server output.
import workflowV1 from '../assets/workflow_minimax_video.json'
import workflowV2 from '../assets/workflow_minimax_video_v2.json'
import workflowLtx from '../assets/workflow_ltx_video.json'
import { MUSIC_NODE_LABELS } from '~/server/utils/musicWorkflow'
import { buildIdentityPrompt } from '~/utils/identityPrompt'
import { applyUserLoraToGraph, resolveGraphLoraNames } from '~/server/utils/loras'
import type { LoraStackItem } from '~/utils/loras'
import {
  isLtxWorkflow,
  isTextToVideo,
  LTX_DISABLED_MESSAGE,
  LTX_NEGATIVE,
  ltxWorkflowEnabled,
  parseVideoWorkflow,
  type VideoWorkflowId
} from '~/utils/videoModels'

export type { VideoWorkflowId }
export { parseVideoWorkflow }

export interface GenerateParams {
  prompt: string
  imageName: string
  width: number
  height: number
  steps: number
  seed: number
  turbo: boolean
  length: number
  cfg: number
  fps: number
  samplerName: string
  scheduler: string
  filenamePrefix?: string
  sound?: boolean
  workflow?: VideoWorkflowId
  duration?: number
  useIdentityRefs?: boolean
  referenceImageNames?: string[]
  loraName?: string
  loraStack?: LoraStackItem[]
  saveLosslessAnchor?: boolean
}

type WorkflowNode = { class_type: string; inputs: Record<string, unknown>; _meta?: { title?: string } }
type WorkflowGraph = Record<string, WorkflowNode>

const PROMPT_CLASSES = new Set([
  'MiniMaxH3ImageToVideo',
  'MiniMaxH3TextToVideo',
  'MiniMaxH3ReferenceToVideo',
  'CLIPTextEncode',
  'CLIPTextEncodeQwen3VL',
  'TextEncodeQwen3VL'
])
const SEED_CLASSES = new Set(['RandomNoise', 'KSampler', 'SamplerMiniMax', 'KSamplerAdvanced'])
const IMAGE_CLASSES = new Set(['LoadImage'])
const LORA_CLASSES = new Set(['LoraLoaderModelOnly', 'LoraLoader'])
const V2_REF_LOADERS = ['140', '141', '142', '143'] as const

function entriesOf(graph: WorkflowGraph, classType: string) {
  return Object.entries(graph).filter(([, node]) => node.class_type === classType)
}

function linkId(value: unknown): string | null {
  return Array.isArray(value) && typeof value[0] === 'string' ? value[0] : null
}

function snap32(value: number) {
  return Math.max(32, Math.round(value / 32) * 32)
}

function titleOf(node: WorkflowNode) {
  return node._meta?.title || ''
}

function templateFor(id: VideoWorkflowId) {
  if (isLtxWorkflow(id)) return workflowLtx as WorkflowGraph
  return (id === 'v2' ? workflowV2 : workflowV1) as WorkflowGraph
}

export function buildWorkflow(params: GenerateParams) {
  const version = parseVideoWorkflow(params.workflow)
  if (isLtxWorkflow(version) && !ltxWorkflowEnabled()) {
    throw new Error(LTX_DISABLED_MESSAGE)
  }
  if (isLtxWorkflow(version)) return buildLtxWorkflow(params, version)

  const graph = structuredClone(templateFor(version === 't2v' ? 'v1' : version === 'v2' ? 'v2' : 'v1'))
  const width = snap32(params.width)
  const height = snap32(params.height)
  const refs = Array.from({ length: 4 }, (_, index) => String(params.referenceImageNames?.[index] || '').trim())
  const extraPictures = refs.flatMap((name, index) => name ? [index + 2] : [])
  const useRefs = version === 'v2' && params.useIdentityRefs === true
  const actionPrompt = params.prompt
  const graphPrompt = useRefs ? buildIdentityPrompt(actionPrompt, extraPictures) : actionPrompt

  for (const [id, node] of Object.entries(graph)) {
    if (IMAGE_CLASSES.has(node.class_type) && 'image' in node.inputs && id === '114') {
      node.inputs.image = params.imageName
    }

    if (PROMPT_CLASSES.has(node.class_type) || node.class_type.includes('Qwen3')) {
      const nodePrompt = promptForNode(node.class_type, useRefs, actionPrompt, graphPrompt)
      if ('prompt' in node.inputs) node.inputs.prompt = nodePrompt
      if ('text' in node.inputs) node.inputs.text = nodePrompt
    }

    if (node.class_type === 'MiniMaxH3ImageToVideo') {
      node.inputs.prompt = useRefs ? '' : actionPrompt
      if (typeof node.inputs.width === 'number') node.inputs.width = width
      if (typeof node.inputs.height === 'number') node.inputs.height = height
      if (typeof node.inputs.length === 'number') node.inputs.length = params.length
    }

    if (node.class_type === 'MiniMaxH3ReferenceToVideo') {
      node.inputs.prompt = useRefs ? graphPrompt : actionPrompt
      if (!Array.isArray(node.inputs.length)) node.inputs.length = ['174', 1]
    }

    if (node.class_type === 'CreateVideo' && 'fps' in node.inputs) {
      node.inputs.fps = params.fps
    }

    if (node.class_type === 'KSamplerSelect' && 'sampler_name' in node.inputs) {
      node.inputs.sampler_name = params.samplerName
    }

    if (node.class_type === 'BasicScheduler' && 'scheduler' in node.inputs) {
      node.inputs.scheduler = params.scheduler
    }

    if ('cfg' in node.inputs && typeof node.inputs.cfg === 'number') {
      node.inputs.cfg = params.cfg
    }
    if ('guidance' in node.inputs && typeof node.inputs.guidance === 'number') {
      node.inputs.guidance = params.cfg
    }

    if (SEED_CLASSES.has(node.class_type)) {
      if ('noise_seed' in node.inputs) node.inputs.noise_seed = params.seed
      if ('seed' in node.inputs) node.inputs.seed = params.seed
    }

    if (LORA_CLASSES.has(node.class_type) && !/identity|ref 1/i.test(titleOf(node))) {
      node.inputs.strength_model = params.turbo ? 1 : 0
    }

    if (node.class_type === 'SaveVideo' && 'filename_prefix' in node.inputs) {
      node.inputs.filename_prefix = params.filenamePrefix || node.inputs.filename_prefix || 'video/MiniMax_H3'
    }

    if (node.class_type === 'PrimitiveBoolean' && /lightning|turbo/i.test(titleOf(node))) {
      node.inputs.value = params.turbo
    }

    if (node.class_type === 'PrimitiveBoolean' && /identity|ref2va|reference/i.test(titleOf(node))) {
      node.inputs.value = useRefs
    }

    if (node.class_type === 'PrimitiveFloat' && /duration/i.test(titleOf(node))) {
      const seconds = typeof params.duration === 'number' && params.duration > 0
        ? params.duration
        : params.length / Math.max(params.fps, 1)
      node.inputs.value = seconds
    }

    if (node.class_type === 'ComfySwitchNode' && typeof node.inputs.switch === 'boolean' && /lightning|turbo|lora/i.test(titleOf(node))) {
      node.inputs.switch = params.turbo
    }

    if (node.class_type === 'BasicScheduler' && typeof node.inputs.steps === 'number') {
      node.inputs.steps = params.steps
    }
  }

  if (version === 'v1' && !isTextToVideo(version) && graph['128']?.class_type === 'ImageScaleToTotalPixels') {
    graph['128'] = {
      class_type: 'ImageScale',
      inputs: {
        upscale_method: 'nearest-exact',
        width,
        height,
        crop: 'center',
        image: ['114', 0]
      },
      _meta: { title: 'Scaling image' }
    }
  }

  if (version === 'v1' || version === 't2v') {
    graph['105:cfg'] = {
      class_type: 'FluxGuidance',
      inputs: {
        guidance: params.cfg,
        conditioning: ['105:104', 0]
      },
      _meta: { title: 'Guidance scale' }
    }
    if (graph['105:16']?.inputs) {
      graph['105:16'].inputs.conditioning = ['105:cfg', 0]
    }
  }

  for (const [, node] of entriesOf(graph, 'ComfySwitchNode')) {
    if (!/step/i.test(titleOf(node))) continue
    const qualityId = linkId(node.inputs.on_false)
    const turboId = linkId(node.inputs.on_true)
    if (qualityId && graph[qualityId]?.class_type === 'PrimitiveInt') {
      graph[qualityId].inputs.value = params.turbo ? 20 : params.steps
    }
    if (turboId && graph[turboId]?.class_type === 'PrimitiveInt') {
      graph[turboId].inputs.value = params.turbo ? params.steps : (version === 'v2' ? 6 : 8)
    }
  }

  if (version === 't2v') applyMinimaxTextToVideo(graph, params, width, height)
  if (version === 'v2') applyV2IdentityPath(graph, params, useRefs, refs, graphPrompt)

  if (params.sound === false) {
    delete graph['105:23']
    delete graph['105:24']
    if (graph['105:91']?.inputs) delete graph['105:91'].inputs.audio
    delete graph['163']
    if (graph['172']?.inputs) delete graph['172'].inputs.audio
    if (graph['105:104']?.inputs) delete graph['105:104'].inputs.audio_vae
  }

  resolveGraphLoraNames(graph, 'video')
  applyUserLoraToGraph(graph, params.loraStack || params.loraName)
  attachLosslessAnchorSave(graph, params, version === 'v2' ? '164' : '105:90')
  return graph
}

function attachLosslessAnchorSave(graph: WorkflowGraph, params: GenerateParams, decodeNodeId: string) {
  if (!params.saveLosslessAnchor) return
  if (!graph[decodeNodeId]) return
  const batchIndex = Math.max(0, Math.floor(Number(params.length) || 1) - 1)
  const prefix = String(params.filenamePrefix || 'video/MiniMax_H3').replace(/\/+$/, '')
  graph.aigen_anchor_pick = {
    class_type: 'ImageFromBatch',
    inputs: {
      image: [decodeNodeId, 0],
      batch_index: batchIndex,
      length: 1
    },
    _meta: { title: 'Last Decoded Frame' }
  }
  graph.aigen_anchor_save = {
    class_type: 'SaveImage',
    inputs: {
      filename_prefix: `${prefix}_last`,
      images: ['aigen_anchor_pick', 0]
    },
    _meta: { title: 'Save Lossless Anchor' }
  }
}

function promptForNode(classType: string, useRefs: boolean, actionPrompt: string, graphPrompt: string) {
  if (classType === 'MiniMaxH3ImageToVideo') return useRefs ? '' : actionPrompt
  if (classType === 'MiniMaxH3TextToVideo') return actionPrompt
  if (classType === 'MiniMaxH3ReferenceToVideo') return useRefs ? graphPrompt : actionPrompt
  return graphPrompt
}

function applyMinimaxTextToVideo(graph: WorkflowGraph, params: GenerateParams, width: number, height: number) {
  const node = graph['105:104']
  if (node) {
    // MiniMax H3 has no TextToVideo class. T2V is ImageToVideo with first_frame disconnected.
    node.class_type = 'MiniMaxH3ImageToVideo'
    node._meta = { title: 'MiniMax H3 Text to Video' }
    delete node.inputs.first_frame
    delete node.inputs.last_frame
    node.inputs.prompt = params.prompt
    node.inputs.width = width
    node.inputs.height = height
  }
  delete graph['114']
  delete graph['119']
  delete graph['120']
  delete graph['127']
  delete graph['128']
}

function buildLtxWorkflow(params: GenerateParams, version: VideoWorkflowId) {
  const graph = structuredClone(workflowLtx as WorkflowGraph)
  const width = snap32(params.width)
  const height = snap32(params.height)
  const textToVideo = isTextToVideo(version)
  if (graph['4']?.inputs) graph['4'].inputs.image = params.imageName
  if (graph['5']?.inputs) {
    graph['5'].inputs.width = width
    graph['5'].inputs.height = height
  }
  if (graph['6']?.inputs) graph['6'].inputs.text = params.prompt
  if (graph['7']?.inputs) graph['7'].inputs.text = LTX_NEGATIVE
  if (graph['8']?.inputs && typeof graph['8'].inputs.length === 'number') graph['8'].inputs.length = params.length
  if (graph['9']?.inputs) graph['9'].inputs.frame_rate = params.fps
  if (graph['10']?.inputs) graph['10'].inputs.cfg = params.cfg
  if (graph['11']?.inputs) graph['11'].inputs.sampler_name = params.samplerName || 'euler'
  if (graph['12']?.inputs) graph['12'].inputs.steps = params.steps
  if (graph['13']?.inputs) graph['13'].inputs.noise_seed = params.seed
  if (graph['16']?.inputs) {
    graph['16'].inputs.frame_rate = params.fps
    graph['16'].inputs.filename_prefix = params.filenamePrefix || 'video/LTX23'
  }
  if (graph['21']?.inputs) graph['21'].inputs.strength_model = 1
  if (textToVideo) {
    delete graph['4']
    delete graph['5']
    graph['8'] = {
      class_type: 'EmptyLTXVLatentVideo',
      inputs: {
        width,
        height,
        length: params.length,
        batch_size: 1
      },
      _meta: { title: 'Empty LTX latent' }
    }
    if (graph['9']?.inputs) {
      graph['9'].inputs.positive = ['6', 0]
      graph['9'].inputs.negative = ['7', 0]
    }
    if (graph['12']?.inputs) graph['12'].inputs.latent = ['8', 0]
    if (graph['14']?.inputs) graph['14'].inputs.latent_image = ['8', 0]
  }
  resolveGraphLoraNames(graph, 'video')
  applyUserLoraToGraph(graph, params.loraStack || params.loraName)
  // LTX VAEDecode is typically node 15
  attachLosslessAnchorSave(graph, params, '15')
  return graph
}

function applyV2IdentityPath(
  graph: WorkflowGraph,
  params: GenerateParams,
  useRefs: boolean,
  refs: string[],
  graphPrompt: string
) {
  if (useRefs && graph['155'] && graph['167'] && graph['168']) {
    graph['167'].inputs.latent_image = ['155', 1]
    graph['168'].inputs.conditioning = ['155', 0]
    graph['155'].inputs.prompt = graphPrompt
    graph['155'].inputs.length = ['174', 1]
    delete graph['173']
    V2_REF_LOADERS.forEach((id, index) => {
      if (!graph[id]) return
      const name = refs[index]
      if (name) {
        graph[id].inputs.image = name
        return
      }
      delete graph[id]
      delete graph['155'].inputs[`ref_images.ref_image_${index + 1}`]
    })
    if (graph['150']?.inputs) graph['150'].inputs.value = true
    return
  }

  delete graph['155']
  delete graph['150']
  for (const id of V2_REF_LOADERS) delete graph[id]
  if (graph['173'] && graph['168']) {
    graph['v2:cfg'] = {
      class_type: 'FluxGuidance',
      inputs: {
        guidance: params.cfg,
        conditioning: ['173', 0]
      },
      _meta: { title: 'Guidance scale' }
    }
    graph['168'].inputs.conditioning = ['v2:cfg', 0]
  }
}

const LABEL_OVERRIDES: Record<string, string> = {
  '114': 'Loading image',
  '128': 'Scaling image',
  '127': 'Reading image size',
  '119': 'Scaling image',
  '120': 'Reading image size',
  '92': 'Saving MP4',
  '140': 'Loading identity ref 1',
  '141': 'Loading identity ref 2',
  '142': 'Loading identity ref 3',
  '143': 'Loading identity ref 4',
  '150': 'Identity references',
  '155': 'Reference to video',
  '172': 'Encoding video',
  '173': 'Image to video',
  '105:6': 'Loading MiniMax H3 model',
  '105:11': 'Loading video VAE',
  '105:13': 'Loading Qwen3-VL encoder',
  '105:15': 'Preparing seed',
  '105:9': 'Building schedule',
  '105:10': 'Decoding video frames',
  '105:14': 'Sampling',
  '105:16': 'Preparing guider',
  '105:cfg': 'Applying guidance',
  'v2:cfg': 'Applying guidance',
  '105:17': 'Selecting sampler',
  '105:23': 'Decoding audio',
  '105:24': 'Loading audio VAE',
  '105:91': 'Encoding video',
  '105:104': 'Encoding prompt & latent',
  '105:107': 'Computing duration',
  '105:111': 'Setting duration',
  'user:lora': 'Applying LoRA',
  '105:121': 'Applying turbo LoRA',
  '105:122': 'Selecting model path',
  '105:123': 'Selecting steps',
  '105:124': 'Quality steps',
  '105:125': 'Turbo steps',
  '105:126': 'Selecting turbo LoRA'
}

function labelsFrom(template: WorkflowGraph) {
  return Object.entries(template).map(([id, node]) => [id, node._meta?.title || `Node ${id}`] as const)
}

export const NODE_LABELS: Record<string, string> = {
  ...Object.fromEntries(labelsFrom(workflowV1 as WorkflowGraph)),
  ...Object.fromEntries(labelsFrom(workflowV2 as WorkflowGraph)),
  ...Object.fromEntries(labelsFrom(workflowLtx as WorkflowGraph)),
  ...MUSIC_NODE_LABELS,
  ...LABEL_OVERRIDES
}

export function isEncodingNode(node: string) {
  return node === '92' || node === '16' || node === '105:91' || node === '172' || /encoding|saving mp4|create video|save mp4/i.test(NODE_LABELS[node] || '')
}