import {sampleKleinSource} from './graphs.mjs'

export function kleinIdentityPlan(request) {
 if(request.engine!=='flux' || ['video','extend'].includes(request.mode)) return null
 const sourceId=request.mode==='generate' ? request.identityStillId || request.imageAId || '' : request.imageAId || request.identityStillId || ''
 const mode=['generate','iterate'].includes(request.mode) ? (sourceId?'edit':'generate') : request.mode
 return {sourceId,mode,denoise:request.lockFace!==false?.45:.65}
}

/** Validate the sampler's actual latent connection, not merely the presence of a reference. */
export function assertKleinSample(graph, sourceId) {
 let link=graph['19']?.inputs?.latent_image
 const visited=new Set()
 while(Array.isArray(link) && !visited.has(String(link[0]))) {
  const id=String(link[0]);visited.add(id)
  const node=graph[id]
  if(node?.class_type==='VAEEncode' && id==='11') return 'source'
  if(node?.class_type?.startsWith('Empty')) {
   if(sourceId) throw new Error('Klein identity requires the source image latent; empty-latent sampling was rejected.')
   return 'empty'
  }
  link=node?.class_type==='SetLatentNoiseMask'?node.inputs.samples:null
 }
 throw new Error('Klein sample latent could not be verified against the source still.')
}

export function applyKleinIdentity(graph, plan) {
 if(plan.sourceId && plan.mode!=='refine') sampleKleinSource(graph,plan.denoise)
 // Retain the masked latent and any already lower refine denoise.
 if(plan.sourceId && plan.mode==='refine') graph['17'].inputs.denoise=Math.min(graph['17'].inputs.denoise,plan.denoise)
 const sampleLatent=assertKleinSample(graph,plan.sourceId)
 return {sampleLatent,sampleDenoise:sampleLatent==='source'?(graph.source_sigmas?.inputs.denoise ?? graph['17']?.inputs.denoise):null}
}