import {kleinIdentityPlan,applyKleinIdentity,assertKleinSample} from '../shared/studio2/klein-identity.mjs'
import {stylePrompt} from '../shared/studio2/styles.mjs'
import test from 'node:test'
import assert from 'node:assert/strict'
import fs from 'node:fs'
import ts from 'typescript'
import {validateRequest,compilePrompt} from '../shared/studio2/contracts.mjs'
import {attachHeroReference,applyResolvedImageSize} from '../shared/studio2/graphs.mjs'
import {extendSelected,changeSource,restoreAsset} from '../shared/studio2/input.mjs'
const template=JSON.parse(fs.readFileSync(new URL('../server/assets/klein_v2_edit.json',import.meta.url)))
const source=fs.readFileSync(new URL('../server/utils/studio2/runner.ts',import.meta.url),'utf8')
const tree=ts.createSourceFile('runner.ts',source,ts.ScriptTarget.Latest,true)
const fn=tree.statements.find(n=>ts.isFunctionDeclaration(n)&&n.name.text==='prepareGraph')
const js=ts.transpileModule(fn.getText(tree),{compilerOptions:{target:ts.ScriptTarget.ES2022}}).outputText
async function prepare(request){
 const uploads=[]
 const scope={stylePrompt,useRuntimeConfig:()=>({comfyFilenamePrefix:'preview'}),upload:async(r,label)=>{uploads.push(label);return label+'.png'},readFileSync:()=>Buffer.from('fixture'),stillPath:()=>'',getStill:(_,id)=>({id}),resolveRequestSize:async()=>{},saveRecord(){},
 buildImageV2Workflow:params=>{const graph=JSON.parse(fs.readFileSync(new URL('../server/assets/klein_v2_'+params.mode+'.json',import.meta.url)));if(graph['1']) graph['1'].inputs.image=params.imageAName;return {graph}},kleinIdentityPlan,applyKleinIdentity,attachHeroReference,applyResolvedImageSize,ensureComfyLoraNames:async()=>{},resolveGraphLoraNames(){}}
 const run=new Function(...Object.keys(scope),`${js};return prepareGraph`)(...Object.values(scope))
 const record={request,owner:'o',id:'job',index:0};return {graph:await run(record),record,uploads}
}
for(const sourceId of ['hero','source','']) test(`Klein edit samples source with hero ${sourceId || 'as sole source'}`,async()=>{
 const q=validateRequest({folderId:'f',mode:'edit',engine:'flux',imageAId:sourceId,identityStillId:'hero',lockFace:true,lockOutfit:false,prompt:'she at a pool'})
 const {graph,record,uploads}=await prepare(q)
 assert.deepEqual(graph['19'].inputs.latent_image,['11',0]);assert.equal(graph['14'],undefined)
 assert.deepEqual(graph['19'].inputs.sigmas,['source_sigmas',1]);assert.equal(graph.source_sigmas.inputs.denoise,.45)
 assert.equal(record.sourceStillId,sourceId||'hero');assert.equal(record.sampleLatent,'source')
 assert.equal(uploads.includes('hero'),sourceId==='source');assert.equal(!!graph.hero_load,sourceId==='source')
 assert.match(q.compiledPrompt,/face, glasses, hair/);assert.match(q.compiledPrompt,/do not restore the hero outfit/)
})
test('image requests omit video sections while form keeps them',()=>{
 const sections={action:'sitting in a restaurant',cameraBeats:[{moves:['hold']}],soundscapeChips:['footsteps'],dialogue:{mode:'whisper',line:'Hello'},musicChip:'low underscore'}
 const q=validateRequest({folderId:'f',mode:'edit',engine:'flux',identityStillId:'hero',promptSections:sections})
 assert.doesNotMatch(q.compiledPrompt,/tripod|hold|soundscape|music|Hello|Picture 1/i)
 assert.deepEqual(q.promptSections,{action:sections.action});assert.equal(sections.cameraBeats.length,1)
 assert.equal(compilePrompt(sections,false,{engine:'krea',identityStillId:'hero'}),sections.action)
})
test('Extend normalizes engine and inherits clip without replacing action or camera',()=>{
 const form={mode:'edit',engine:'flux',imageAId:'still',promptSections:{action:'New action',cameraBeats:[{moves:['hold']}]},settings:{loraStack:[]}}
 const clip={id:'clip',kind:'clip',studio2:{engine:'minimax',identityStillId:'hero',lockOutfit:false,promptSections:{action:'Old action'},settings:{duration:8,aspect:'1:1',width:960,height:960,fps:24,turbo:false,steps:8,cfg:1.5,loraStack:[{name:'a',strengthModel:.6,strengthClip:.3}]}}}
 const next=extendSelected(form,clip);assert.equal(next.engine,'minimax');assert.equal(next.startClipId,'clip');assert.deepEqual(next.settings,clip.studio2.settings)
 assert.deepEqual(next.promptSections,form.promptSections);assert.equal(next.lockOutfit,false)
 assert.equal(changeSource({...next,imageAId:'new-still'}).startClipId,'clip');assert.equal(changeSource(next).mode,'extend')
 assert.deepEqual(restoreAsset(form,clip).promptSections,form.promptSections)
 assert.equal(extendSelected({...form,startClipId:'clip'},null).startClipId,'clip')
 assert.throws(()=>validateRequest({folderId:'f',mode:'extend',engine:'minimax',identityStillId:'hero',prompt:'Go'}),/Select the clip in the family strip or Previous clip/)
})

for (const imageAId of ['', 'hero', 'other']) test('Generate samples hero exclusively: '+imageAId,async()=>{
 const q=validateRequest({folderId:'f',mode:'generate',engine:'flux',imageAId,identityStillId:'hero',lockFace:true,lockOutfit:false,prompt:'she at a pool',settings:{loraStack:[]}})
 const {graph,record,uploads}=await prepare(q)
 assert.equal(record.graphId,'klein_v2_edit.json')
 assert.equal(record.sourceStillId,'hero')
 assert.equal(record.sampleLatent,'source')
 assert.equal(record.sampleDenoise,.45)
 assert.deepEqual(graph['19'].inputs.latent_image,['11',0])
 assert.deepEqual(uploads,['source'])
 assert.equal(graph.hero_load,undefined)
 assert.deepEqual(q.settings.loraStack,[])
 assert.doesNotMatch(q.compiledPrompt,/Picture 1/)
})
test('unlocked source uses .65; text-only uses empty latent',async()=>{
 const q=validateRequest({folderId:'f',mode:'edit',engine:'flux',imageAId:'source',lockFace:false,prompt:'pool'})
 assert.equal((await prepare(q)).record.sampleDenoise,.65)
 const empty=validateRequest({folderId:'f',mode:'generate',engine:'flux',prompt:'pool'})
 assert.equal((await prepare(empty)).record.sampleLatent,'empty')
})
test('reject empty sample with source and exclude other engines',()=>{
 assert.throws(()=>assertKleinSample(template,'hero'),/empty-latent sampling was rejected/)
 assert.equal(kleinIdentityPlan({engine:'krea',mode:'generate',identityStillId:'hero'}),null)
 assert.equal(kleinIdentityPlan({engine:'minimax',mode:'extend'}),null)
})