import template from '../assets/workflow_qwen_edit.json'

export const DEFAULT_EDIT_NEGATIVE = 'blurry, low quality, deformed, bad anatomy, watermark, text, extra fingers'

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

export interface EditParams {
  prompt: string
  negative?: string
  imageName: string
  steps: number
  seed: number
  cfg: number
  denoise: number
  samplerName?: string
  scheduler?: string
  filenamePrefix?: string
}

export function buildEditWorkflow(params: EditParams) {
  const graph = structuredClone(template) as WorkflowGraph
  const loadImage = nodeByClass(graph, 'LoadImage') || graph['4']
  const qwenPrompt = nodeByClass(graph, 'TextEncodeQwenImageEdit') || graph['17']
  const sampler = nodeByClass(graph, 'KSampler') || graph['8']
  const save = nodeByClass(graph, 'SaveImage') || graph['10']
  if (loadImage) loadImage.inputs.image = params.imageName
  if (qwenPrompt) {
    if ('prompt' in qwenPrompt.inputs) qwenPrompt.inputs.prompt = params.prompt
    if ('text' in qwenPrompt.inputs) qwenPrompt.inputs.text = params.prompt
  }
  if (sampler) {
    sampler.inputs.seed = params.seed
    sampler.inputs.steps = params.steps
    sampler.inputs.cfg = params.cfg
    sampler.inputs.denoise = params.denoise
    sampler.inputs.sampler_name = params.samplerName || 'euler'
    sampler.inputs.scheduler = params.scheduler || 'simple'
  }
  if (save && params.filenamePrefix) save.inputs.filename_prefix = params.filenamePrefix
  return graph
}

function nodeByClass(graph: WorkflowGraph, classType: string, index = 0) {
  const matches = Object.values(graph).filter(node => node?.class_type === classType)
  return matches[index]
}

export const EDIT_NODE_LABELS: Record<string, string> = {
  '1': 'Loading Qwen Edit UNET',
  '11': 'Applying Lightning LoRA',
  '2': 'Loading Qwen CLIP',
  '3': 'Loading Qwen VAE',
  '4': 'Loading input still',
  '17': 'Encoding Qwen edit prompt',
  '19': 'Preparing negative',
  '7': 'Encoding still',
  '8': 'Sampling edit',
  '9': 'Decoding still',
  '10': 'Saving still'
}