'use client';
import {useCallback,useRef,useState} from 'react';
import {type Clip,type Project,type Range,type Selection,addTrack,clipEnd,clipsIn,muteRange,projectEnd,soleClip,uid} from '@/lib/studio/model';
import {type Player,renderMix,sliceBuffer,spliceBuffer} from '@/lib/studio/audio';
import {SignedOut,callEngine,fetchResult,wavFile} from '@/lib/studio/engine';
import type {Analysis,BandResult,CoverResult,Download,LyricsResult,PartResult,ScoreRead,SongResult,StemsResult} from '@/lib/types';
import {type Model,modelLabel} from '@/lib/studio/models';

export type Job={id:string;label:string;status:string;start:number;end:number;error?:string};

/**
 * A score the user read from a selection and edited: the engine plays this
 * instead of reading the audio again. `clock` is when each bar starts in the
 * selection, which is what the result is locked onto.
 */
export type ScoreDraft={abc:string;clock:number[];beats?:number[];edits:{transpose?:number;bpm?:number;chords?:Record<number,string[]>}};

function addScore(form:FormData,score?:ScoreDraft){
 if(!score)return;
 form.append('abc',score.abc);form.append('clock',JSON.stringify({bars:score.clock,beats:score.beats||[]}));form.append('edits',JSON.stringify(score.edits));
}

/** What an AI action works on: the clips it reads and the stretch of time it covers. */
export type Target={start:number;end:number;clips:Clip[];tracks:string[];range:Range|null;
 /** The single clip that covers the whole target, so a result can be spliced straight into it. */
 sole:Clip|null};

export function targetOf(project:Project,selection:Selection):Target|null{
 const range=selection.range;
 if(range&&range.end-range.start>0.05){
  const clips=clipsIn(project,range).filter(c=>!c.muted);
  if(!clips.length)return null;
  // A range drawn past the audio would send silence the engine then has to fill.
  const start=Math.max(range.start,Math.min(...clips.map(c=>c.start)));
  const end=Math.min(range.end,Math.max(...clips.map(clipEnd)));
  const clamped={...range,start,end};
  return {start,end,clips,tracks:range.tracks,range:clamped,sole:soleClip(project,clamped)};
 }
 const clips=project.clips.filter(c=>selection.clips.includes(c.id));
 if(!clips.length)return null;
 return {start:Math.min(...clips.map(c=>c.start)),end:Math.max(...clips.map(clipEnd)),clips,
  tracks:[...new Set(clips.map(c=>c.trackId))],range:null,sole:clips.length===1?clips[0]:null};
}

type Deps={
 player:Player;buffers:Map<string,AudioBuffer>;
 project:()=>Project;playhead:()=>number;
 commit:(fn:(p:Project)=>Project)=>void;
 signedOut:()=>void;
};
type Placed={name:string;buffer:AudioBuffer;muted?:boolean};
const MAX=240;
/** Score jobs (YuE2 covers and bands, reading scores and words) take a whole song. */
const SCORE_MAX=300;

export function useJobs({player,buffers,project,playhead,commit,signedOut}:Deps){
 const [jobs,setJobs]=useState<Job[]>([]);
 const controllers=useRef(new Map<string,AbortController>());

 const store=useCallback((buffer:AudioBuffer)=>{const id=uid();buffers.set(id,buffer);return id;},[buffers]);
 const load=useCallback(async(file:Download|undefined,signal:AbortSignal)=>{
  if(!file)throw new Error('The engine did not return that audio.');
  return player.decode(await fetchResult(file,signal));
 },[player]);

 /**
  * Run one engine request as a job. Jobs are independent: the user keeps
  * editing while the GPU works, so each result is placed against the project
  * as it is when the job finishes, not as it was when it started.
  */
 const run=useCallback((label:string,start:number,end:number,work:(signal:AbortSignal,status:(t:string)=>void)=>Promise<void>)=>{
  const id=uid(),controller=new AbortController();
  controllers.current.set(id,controller);
  setJobs(list=>[...list,{id,label,status:'Starting…',start,end}]);
  const status=(text:string)=>setJobs(list=>list.map(j=>j.id===id?{...j,status:text}:j));
  work(controller.signal,status).then(()=>{
   setJobs(list=>list.filter(j=>j.id!==id));
  },error=>{
   if(controller.signal.aborted){setJobs(list=>list.filter(j=>j.id!==id));return;}
   if(error instanceof SignedOut)signedOut();
   setJobs(list=>list.map(j=>j.id===id?{...j,error:error instanceof Error?error.message:'That could not finish.'}:j));
  }).finally(()=>controllers.current.delete(id));
 },[signedOut]);

 const cancel=useCallback((id:string)=>{controllers.current.get(id)?.abort();setJobs(list=>list.filter(j=>j.id!==id));},[]);

 /** Put results on new tracks below the source and, when they replace it, mute what they replace. */
 const addBelow=useCallback((target:Target|null,at:number,items:Placed[],replace:boolean)=>{
  const made=items.map(item=>({...item,id:store(item.buffer)}));
  commit(p=>{
   let next=p;
   let after=target?[...p.tracks].reverse().find(t=>target.tracks.includes(t.id))?.id:undefined;
   for(const item of made){
    const [withTrack,track]=addTrack(next,item.name,after);after=track.id;
    next={...withTrack,clips:[...withTrack.clips,{id:uid(),trackId:track.id,buffer:item.id,start:at,offset:0,
     duration:item.buffer.duration,name:item.name,muted:item.muted}]};
   }
   if(replace&&target){
    next=target.range?muteRange(next,target.range)
     :{...next,clips:next.clips.map(c=>target.clips.some(t=>t.id===c.id)?{...c,muted:true}:c)};
   }
   return next;
  });
 },[commit,store]);

 /**
  * Place a result that covers exactly the target: splice it into the one clip
  * it came from when there is one, otherwise lay it on a new track.
  */
 const place=useCallback((target:Target,result:AudioBuffer,name:string)=>{
  const length=target.end-target.start;
  const piece=result.duration>length+0.01?sliceBuffer(result,0,length):result;
  const sole=target.sole&&project().clips.find(c=>c.id===target.sole!.id);
  const source=sole&&buffers.get(sole.buffer);
  if(sole&&source){
   const at=target.sole!.offset+(target.start-target.sole!.start);
   const id=store(spliceBuffer(source,at,piece));
   commit(p=>({...p,clips:p.clips.map(c=>c.id===sole.id?{...c,buffer:id}:c)}));
   return;
  }
  addBelow(target,target.start,[{name,buffer:piece}],true);
 },[addBelow,buffers,commit,project,store]);

 /** Audio for a target. One clip goes raw; several are mixed as the user hears them against each other. */
 const render=useCallback((target:Target,start=target.start,end=target.end)=>
  renderMix(project(),buffers,{start,end,clips:target.clips.map(c=>c.id),mixer:target.clips.length>1}),[buffers,project]);

 const check=(target:Target|null,min=3,max=MAX)=>{
  if(!target)throw new Error('Select a clip, or drag across one, first.');
  const length=target.end-target.start;
  if(length<min)throw new Error(`Select at least ${min} seconds.`);
  // Name what was measured: several clips count from the first start to the last end.
  if(length>max)throw new Error(`Your selection covers ${Math.floor(length/60)}:${String(Math.floor(length%60)).padStart(2,'0')}`
   +`${target.clips.length>1?` across ${target.clips.length} clips`:''}. Select ${max/60} minutes or less at a time.`);
  return target;
 };

 const restyle=useCallback((selection:Selection,options:{style:string;change:number;lyrics:string;model:Model})=>{
  const target=check(targetOf(project(),selection));
  run(`Restyle · ${modelLabel(options.model)}`,target.start,target.end,async(signal,status)=>{
   const form=new FormData();
   form.append('audio',wavFile(await render(target)));
   form.append('model',options.model);form.append('style',options.style);form.append('change',String(options.change));form.append('lyrics',options.lyrics);
   const result=await callEngine<{audio?:Download}>('/api/remix',form,{signal,status});
   status('Placing the result…');
   place(target,await load(result.audio,signal),'Restyled');
  });
 },[load,place,project,render,run]);

 const rewrite=useCallback((selection:Selection,options:{style:string;variance:number;lyrics:string;model:Model})=>{
  const target=check(targetOf(project(),selection),1);
  // Send some audio either side so the model hears what the new part must join.
  const from=Math.min(...target.clips.map(c=>c.start)),to=Math.max(...target.clips.map(clipEnd));
  let lo=Math.max(from,target.start-8),hi=Math.min(to,target.end+8);
  if(hi-lo<3){lo=Math.max(0,Math.min(lo,target.end-3));hi=Math.max(hi,lo+3);}
  if(hi-lo>MAX)throw new Error('Select four minutes or less at a time.');
  run(`Rewrite · ${modelLabel(options.model)}`,target.start,target.end,async(signal,status)=>{
   const form=new FormData();
   form.append('audio',wavFile(await render(target,lo,hi)));
   form.append('model',options.model);form.append('style',options.style);form.append('lyrics',options.lyrics);
   form.append('start',String(target.start-lo));form.append('end',String(target.end-lo));
   form.append('variance',String(options.variance));
   const result=await callEngine<{audio?:Download}>('/api/repaint',form,{signal,status});
   status('Placing the result…');
   const full=await load(result.audio,signal);
   place(target,sliceBuffer(full,target.start-lo,target.end-lo),'Rewritten');
  });
 },[load,place,project,render,run]);

 /** Adopt the tempo, key and chords the engine measured so the grid and later parts line up. */
 const learn=useCallback((analysis:Analysis|undefined,offset:number)=>{
  if(!analysis)return;
  commit(p=>({...p,
   bpm:analysis.tempo>0?Math.round(analysis.tempo*10)/10:p.bpm,
   key:analysis.key||p.key,
   chords:analysis.chords.length?analysis.chords.map(c=>({time:c.time+offset,name:c.name})):p.chords}));
 },[commit]);

 const band=useCallback((selection:Selection,options:{style:string;separate:string;strength:number;split:boolean;useTempo:boolean;model:Model;score?:ScoreDraft})=>{
  // YuE2 and ACE-Step 1.5 bands take a whole song; ACE-Step 1 stays at four minutes.
  const target=check(targetOf(project(),selection),3,options.model==='ace1'?MAX:SCORE_MAX);
  const bpm=options.useTempo?project().bpm:0;
  run(`Band · ${modelLabel(options.model)}`,target.start,target.end,async(signal,status)=>{
   const form=new FormData();
   form.append('audio',wavFile(await render(target)));
   form.append('model',options.model);form.append('style',options.style);form.append('separate',options.separate);
   form.append('strength',String(options.strength));form.append('split',String(options.split));form.append('bpm',String(bpm));
   if(options.model==='yue2')addScore(form,options.score);
   const result=await callEngine<BandResult>('/api/accompany',form,{signal,status});
   status('Laying out your band…');
   const items:Placed[]=[];
   const lifted=!!result.analysis?.separated&&!!result.vocal;
   if(lifted){
    items.push({name:'Lifted vocal',buffer:await load(result.vocal,signal)});
    if(result.background)items.push({name:'Original backing',buffer:await load(result.background,signal),muted:true});
   }
   if(result.drums&&result.bass&&result.other){
    // Named apart from stem splits, which also produce drums and bass tracks.
    items.push({name:'Band drums',buffer:await load(result.drums,signal)});
    items.push({name:'Band bass',buffer:await load(result.bass,signal)});
    items.push({name:'Band keys & guitars',buffer:await load(result.other,signal)});
   }else items.push({name:'Band',buffer:await load(result.backing,signal)});
   learn(result.analysis,target.start);
   addBelow(target,target.start,items,lifted);
  });
 },[addBelow,learn,load,project,render,run]);

 /**
  * Perform the selection's tune in a new style with YuE2. The result comes
  * back locked to the selection's bars and exactly as long, so it goes on a
  * new track in place and the source is muted.
  */
 const cover=useCallback((selection:Selection,options:{style:string;lyrics:string;voice:'sing'|'play';harmony:'keep'|'free';score?:ScoreDraft})=>{
  const target=check(targetOf(project(),selection),3,SCORE_MAX);
  run('Cover · YuE2',target.start,target.end,async(signal,status)=>{
   const form=new FormData();
   form.append('audio',wavFile(await render(target)));
   form.append('model','yue2');form.append('style',options.style);form.append('lyrics',options.lyrics);
   form.append('voice',options.voice);form.append('harmony',options.harmony);
   addScore(form,options.score);
   const result=await callEngine<CoverResult>('/api/cover',form,{signal,status});
   status('Placing the cover…');
   const buffer=await load(result.audio,signal),length=target.end-target.start;
   addBelow(target,target.start,[{name:options.voice==='sing'?'Cover':'Instrumental cover',
    buffer:buffer.duration>length+0.01?sliceBuffer(buffer,0,length):buffer}],true);
  });
 },[addBelow,load,project,render,run]);

 /** Read the selection's bars, chords, key and (optionally) words, and teach them to the project. */
 const readScore=useCallback((selection:Selection,options:{words:boolean},done:(read:ScoreRead,target:Target)=>void)=>{
  const target=check(targetOf(project(),selection),3,SCORE_MAX);
  run('Reading the score',target.start,target.end,async(signal,status)=>{
   const form=new FormData();
   form.append('audio',wavFile(await render(target)));
   form.append('words',String(options.words));
   const read=await callEngine<ScoreRead>('/api/score',form,{signal,status});
   const summary=read.summary;
   if(!summary||!read.abc)throw new Error('No score came back. Try a longer stretch.');
   // The score repeats the sounding chord at every bar; the lane marks only the changes.
   const chords=summary.bars.flatMap(bar=>bar.chords.map((name,i)=>
    ({time:bar.start+i*(bar.end-bar.start)/bar.chords.length,name})))
    .filter((chord,i,all)=>i===0||all[i-1].name!==chord.name);
   learn({tempo:summary.bpm,key:summary.keyName,chords},target.start);
   done(read,target);
  });
 },[learn,project,render,run]);

 /** The words sung in the selection, grouped by section. */
 const words=useCallback((selection:Selection,done:(lyrics:string)=>void)=>{
  const target=check(targetOf(project(),selection),3,SCORE_MAX);
  run('Reading the words',target.start,target.end,async(signal,status)=>{
   const form=new FormData();
   form.append('audio',wavFile(await render(target)));
   const result=await callEngine<LyricsResult>('/api/lyrics',form,{signal,status});
   if(!result.lyrics.trim())throw new Error('No words were found in this stretch.');
   done(result.lyrics);
  });
 },[project,render,run]);

 const stems=useCallback((selection:Selection)=>{
  const target=check(targetOf(project(),selection));
  run('Stems',target.start,target.end,async(signal,status)=>{
   const form=new FormData();
   form.append('audio',wavFile(await render(target)));
   const result=await callEngine<StemsResult>('/api/split',form,{signal,status});
   status('Laying out the stems…');
   addBelow(target,target.start,[
    {name:'Vocals',buffer:await load(result.vocals,signal)},{name:'Drums',buffer:await load(result.drums,signal)},
    {name:'Bass',buffer:await load(result.bass,signal)},{name:'Other',buffer:await load(result.other,signal)}],true);
  });
 },[addBelow,load,project,render,run]);

 const extend=useCallback((selection:Selection,options:{style:string;seconds:number;lyrics:string;model:Model})=>{
  const target=check(targetOf(project(),selection));
  if(target.range||!target.sole)throw new Error('Select one clip to extend.');
  if(target.end-target.start>200)throw new Error('Extend works on clips up to 200 seconds.');
  const clip=target.sole;
  run(`Extend · ${modelLabel(options.model)}`,clipEnd(clip),clipEnd(clip)+options.seconds,async(signal,status)=>{
   const form=new FormData();
   form.append('audio',wavFile(await render(target)));
   form.append('model',options.model);form.append('style',options.style);form.append('seconds',String(options.seconds));form.append('lyrics',options.lyrics);
   const result=await callEngine<{audio?:Download}>('/api/extend',form,{signal,status});
   status('Placing the result…');
   const full=await load(result.audio,signal);
   const now=project().clips.find(c=>c.id===clip.id),source=now&&buffers.get(now.buffer);
   if(!now||!source){addBelow(target,clip.start,[{name:'Extended',buffer:full}],false);return;}
   // Keep the original samples up to the join and crossfade into the model's continuation.
   const join=Math.max(0.5,clip.duration-0.3);
   const kept=sliceBuffer(source,now.offset,now.offset+now.duration);
   const joined=spliceBuffer(kept,join,sliceBuffer(full,join,full.duration),0.3);
   const id=store(joined);
   commit(p=>({...p,clips:p.clips.map(c=>c.id===clip.id?{...c,buffer:id,offset:0,duration:joined.duration}:c)}));
  });
 },[addBelow,buffers,commit,load,project,render,run,store]);

 const part=useCallback((options:{instrument:string;label:string;style:string;follow:boolean;duration:number;strength:number;model:Model})=>{
  const current=project();
  const end=projectEnd(current);
  const follow=options.follow&&current.clips.some(c=>!c.muted);
  if(follow&&end<3)throw new Error('The project is too short to follow. Add at least three seconds.');
  // A part follows at most four minutes: from the playhead, or from the top
  // when the playhead sits too near the end to leave three seconds to follow.
  const cursor=playhead();
  const at=follow?(end>MAX&&cursor<end-3?cursor:0):cursor;
  const until=follow?Math.min(end,at+MAX):at+options.duration;
  run(`${options.label} · ${modelLabel(options.model)}`,at,until,async(signal,status)=>{
   const form=new FormData();
   if(follow)form.append('audio',wavFile(await renderMix(current,buffers,{start:at,end:until})));
   form.append('model',options.model);form.append('instrument',options.instrument);form.append('style',options.style);
   // Once a take has been analysed its tempo is known, and every part should share that grid.
   form.append('bpm',String(current.chords.length||!follow?current.bpm:0));
   form.append('key',current.key);form.append('duration',String(options.duration));
   form.append('strength',String(options.strength));
   const result=await callEngine<PartResult>('/api/part',form,{signal,status});
   status('Placing the part…');
   const buffer=await load(result.audio,signal);
   if(!project().chords.length)learn(result.analysis,at);
   addBelow(null,at,[{name:options.label,buffer}],false);
  });
 },[addBelow,buffers,learn,load,playhead,project,run]);

 const generate=useCallback((options:{kind:'instrumental'|'song';style:string;lyrics:string;duration:number;model:Model})=>{
  if(!options.style.trim())throw new Error('Describe the style you want.');
  if(options.kind==='song'&&!options.lyrics.trim())throw new Error('Add lyrics, marked with [verse] and [chorus].');
  const at=playhead(),current=project();
  const style=options.kind==='instrumental'&&current.clips.length
   ?[options.style,`${Math.round(current.bpm)} bpm`,current.key&&`key of ${current.key}`].filter(Boolean).join(', '):options.style;
  run(`${options.kind==='song'?'Song':'Instrumental'} · ${modelLabel(options.model)}`,at,at+(options.kind==='song'?60:options.duration),async(signal,status)=>{
   const result=options.kind==='song'
    ?await callEngine<SongResult>('/api/generate',{style,lyrics:options.lyrics,cot:'full',model:options.model},{signal,status})
    :await callEngine<{audio?:Download}>('/api/instrumental',{style,lyrics:'',duration:options.duration,model:options.model},{signal,status});
   status('Placing the result…');
   addBelow(null,at,[{name:options.kind==='song'?'Song':'Instrumental',buffer:await load(result.audio,signal)}],false);
  });
 },[addBelow,load,playhead,project,run]);

 return {jobs,cancel,restyle,rewrite,band,cover,readScore,words,stems,extend,part,generate,
  dismiss:(id:string)=>setJobs(list=>list.filter(j=>j.id!==id))};
}