// 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 { buildIdentityPrompt } from '~/utils/identityPrompt'

export type VideoWorkflowId = 'v1' | 'v2'

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[]
}

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

const PROMPT_CLASSES = new Set([
  'MiniMaxH3ImageToVideo',
  '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) {
  return (id === 'v2' ? workflowV2 : workflowV1) as WorkflowGraph
}

export function parseVideoWorkflow(raw: unknown): VideoWorkflowId {
  return String(raw || '').trim() === 'v2' ? 'v2' : 'v1'
}

export function buildWorkflow(params: GenerateParams) {
  const version = params.workflow === 'v2' ? 'v2' : 'v1'
  const graph = structuredClone(templateFor(version))
  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 graphPrompt = useRefs ? buildIdentityPrompt(params.prompt, extraPictures) : params.prompt

  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')) {
      if ('prompt' in node.inputs) node.inputs.prompt = graphPrompt
      if ('text' in node.inputs) node.inputs.text = graphPrompt
    }

    if (node.class_type === 'MiniMaxH3ImageToVideo') {
      node.inputs.prompt = graphPrompt
      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 = graphPrompt
      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' && 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') {
    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 : 8
    }
  }

  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
  }

  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']
    if (graph['169']?.inputs) {
      graph['169'].inputs.unet_name = 'minimax_h3_ref2va_pruned_int8_convrot.safetensors'
    }
    if (params.turbo && graph['176']?.inputs) {
      graph['176'].inputs.lora_name = 'minimax_h3_ref2v_turbo_4step_v0.1_comfyui_bf16.safetensors'
    }
    V2_REF_LOADERS.forEach((id, index) => {
      const name = refs[index]
      if (name && graph[id]) {
        graph[id].inputs.image = name
        return
      }
      delete graph[id]
      if (graph['155']?.inputs) 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',
  '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)),
  ...LABEL_OVERRIDES
}

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