import { init, defaultDevice, getWebGPUDevice } from '@jax-js/jax';
import { Trainer, tokenize, batch, paragraphBatch, rng } from './model.js';
import {validateExploration} from '../public/exploration.js';
let exploring=false;
let model, tok, config, random, running=false, busy=false, stopped=false, elapsed=0, request,initialSampleComplete=false;
const post=(type,data={})=>self.postMessage({type,...data});
const yieldTask=()=>new Promise(r=>setTimeout(r,0));
function fail(e){ running=false; stopped=true; post('error',{message:e?.message||String(e)}); }
async function sample(){
 const step=model.step;post('sampling',{step});
 const text=await model.generate(request.prefix,tok,{length:request.sampleLength??80,temperature:request.temperature,seed:123,cancel:()=>!running||stopped,onToken:text=>{if([...text].length%8===0)post('sample-progress',{step,prefix:request.prefix,text});}});
 if(step===0)initialSampleComplete=[...text].length===(request.sampleLength??80);
 post('sample',{step,prefix:request.prefix,text});
}
async function loop(){
 if(busy||stopped)return;busy=true;
 try {
  if(running&&model.step===0&&!initialSampleComplete)await sample();
  while(running&&!stopped){
   const start=performance.now();
   const value=await model.train(...(request.sampling==='paragraphs'?paragraphBatch(tok.ids,config,random,tok.paragraphStarts):batch(tok.ids,config,random)),request.rate*(request.schedule==='cooldown'?Math.min(1,(1-model.step/request.maxSteps)/.7):1));
   const ms=performance.now()-start;elapsed+=ms;
   post('step',{step:model.step,loss:value,ms,elapsed,tokens:model.step*config.batch*config.context});
   if(model.step%4===0)await yieldTask();
   if(running&&(model.step%request.sampleEvery===0||model.step>=request.maxSteps)) await sample();
   if(model.step>=request.maxSteps){running=false;post('complete');}
  }
 }catch(e){fail(e);}finally{busy=false;if(!stopped)post('paused',{step:model.step});}
}
self.onmessage=async({data})=>{
 if(data.type==='cancel-explore'){exploring=false;return;}
 if(data.type==='explore'){
  if(!model||busy||running||stopped||model.step<request.maxSteps){post('explore-error',{id:data.id,message:'Wait for training or the current generation to finish.'});return;}
  busy=true;const started=performance.now();
  try{
   const options=validateExploration(data,request.text);exploring=true;
   const text=await model.generate(options.prefix,tok,{...options,cancel:()=>!exploring||stopped,onToken:text=>{if([...text].length%8===0)post('explore-progress',{id:data.id,text});}});
   post('explore-result',{id:data.id,...options,text,step:model.step,ms:performance.now()-started,cancelled:!exploring});
  }catch(e){post('explore-error',{id:data.id,message:e.message||String(e)});}
  finally{exploring=false;busy=false;}
  return;
 }
 if(data.type==='pause'){running=false;return;}
 if(data.type==='resume'){if(!model||stopped||busy||model.step>=request.maxSteps)return;running=true;post('running');void loop();return;}
 if(data.type!=='start'||busy||model)return;
 busy=true;
 try {
  const loadStart=performance.now();request=data;tok=tokenize(data.text);tok.encode(data.prefix);
  if(!data.prefix.length)throw new Error('Enter a prompt.');
  if(tok.ids.length<129||tok.ids.length>100000)throw new Error('The dataset must contain between 129 and 100,000 characters.');
  post('status',{message:'Initializing the backend…'});
  let available=[];let fallback='';
  if(data.backend!=='wasm') {
   try{available=await init('webgpu');}catch(e){fallback=e.message;}
  }
  const backend=available.includes('webgpu')?'webgpu':'wasm';
  if(data.backend==='webgpu'&&backend!=='webgpu')throw new Error(self.isSecureContext?'WebGPU unavailable. Select Auto or WASM.':'WebGPU unavailable over HTTP: open the demo over HTTPS or select Auto / WASM.');
  if(backend==='wasm')await init('wasm');
  defaultDevice(backend);
  if(backend==='webgpu')getWebGPUDevice().lost.then(info=>{if(!stopped)fail(new Error(`GPU device lost: ${info.message}. Reset in WASM mode.`));});
  const large=data.size==='standard';
  config={architecture:['nanogpt','rotary'].includes(data.architecture)?data.architecture:undefined,fastSampling:data.fastSampling===true&&backend==='wasm',vocab:tok.chars.length,dim:large?64:32,layers:large?2:1,heads:4,context:[32,64,128,256].includes(data.context)?data.context:32,batch:[4,8].includes(data.batch)?data.batch:4,beta2:data.optimizer==='tuned'?.95:.999,weightDecay:data.optimizer==='tuned'?.01:0};
  if(tok.ids.length<=config.context)throw new Error('The dataset must be longer than the selected context.');
  if(data.sampling==='paragraphs'&&tok.paragraphStarts.length<2)throw new Error('This recipe requires at least two paragraphs separated by a blank line.');
  random=rng(99);model=new Trainer(config,42);
  const params=Object.values(model.p).reduce((s,p)=>s+p.size,0);
  post('ready',{loadMs:performance.now()-loadStart,backend,params,vocab:tok.chars.length,config,fallback:backend==='wasm'&&data.backend==='auto'?(self.isSecureContext?'WebGPU unavailable; using WASM on CPU. ':'HTTP connection: using WASM on CPU. WebGPU requires HTTPS. ')+fallback:''});
  running=true;post('status',{message:'First generation and compilation…'});await sample();
  busy=false;if(running){post('running');void loop();}else post('paused',{step:0});
 }catch(e){busy=false;fail(e);}
};