import { Wllama } from '@wllama/wllama';
import wasmURL from '@wllama/wllama/esm/wasm/wllama.wasm?url';
import { createSHA256 } from 'hash-wasm';
import { assertModel, recentMessages } from './model-utils.js';

const folderName = 'cabin-local-models-v1';
const quiet = {log(){},debug(){},warn(...v){console.warn(...v)},error(...v){console.error(...v)}};
const abortCheck = signal => {if (signal?.aborted) throw new DOMException('Operation cancelled', 'AbortError');};
async function directory() {
  if (!navigator.storage?.getDirectory) throw new Error('Model storage is unavailable. Uncheck “Keep the weights” to try in memory, or use a recent browser outside private browsing.');
  return (await navigator.storage.getDirectory()).getDirectoryHandle(folderName,{create:true});
}
export async function cachedModels(models) {
  const ids=[];
  try {const dir=await directory();for(const model of models){try{const f=await (await dir.getFileHandle(model.id)).getFile();if(f.size===model.bytes)ids.push(model.id);}catch{}}}catch{}
  return ids;
}
export async function deleteModel(model) {
  const dir=await directory();try{await dir.removeEntry(model.id);}catch(e){if(e.name!=='NotFoundError')throw e;}
}
export class LocalEngine {
  constructor() { this.runtime=null; this.model=null; this.history=[]; }
  async unload() {
    const old=this.runtime;this.runtime=null;this.model=null;this.history=[];
    if(old)await old.exit();
  }
  async load(model, signal, onProgress, persist=true) {
    await this.unload(); abortCheck(signal);
    const dir=persist?await directory():null; let file;
    if(dir)try {file=await(await dir.getFileHandle(model.id)).getFile();}catch{}
    const hash=await createSHA256();hash.init();
    if(file?.size===model.bytes){
      onProgress({phase:'verify',loaded:0,total:model.bytes});
      const reader=file.stream().getReader();let loaded=0;
      try {while(true){abortCheck(signal);const {done,value}=await reader.read();if(done)break;hash.update(value);loaded+=value.byteLength;onProgress({phase:'verify',loaded,total:model.bytes});}}
      finally{await reader.cancel();}
      try {assertModel(model,file.size,hash.digest());}catch(e){await dir.removeEntry(model.id);throw e;}
    } else {
      const estimate=await navigator.storage?.estimate?.().catch(()=>({}))||{};
      if(persist&&estimate.quota && estimate.quota-estimate.usage < model.bytes*1.1)throw new Error('Not enough storage. Uncheck “Keep the weights” to try in memory, or free some storage.');
      const response=await fetch(model.downloadURL,{signal,credentials:'omit'});
      if(!response.ok || !response.body)throw new Error(`Download unavailable (HTTP ${response.status}). Try again.`);
      let handle,writer;const chunks=[];
      if(dir){try{handle=await dir.getFileHandle(model.id,{create:true});writer=await handle.createWritable();}catch(e){await response.body.cancel().catch(()=>{});await dir.removeEntry(model.id).catch(()=>{});throw e;}}
      const reader=response.body.getReader();let loaded=0;
      try {
        while(true){abortCheck(signal);const {done,value}=await reader.read();if(done)break;loaded+=value.byteLength;if(loaded>model.bytes)throw new Error('Unexpected model file size.');hash.update(value);if(writer)await writer.write(value);else chunks.push(value);onProgress({phase:'download',loaded,total:model.bytes});}
        assertModel(model,loaded,hash.digest());if(writer){await writer.close();file=await handle.getFile();}else{file=new Blob(chunks);chunks.length=0;}
      }catch(e){if(writer)await writer.abort().catch(()=>{});if(dir)await dir.removeEntry(model.id).catch(()=>{});throw e;}
      finally{await reader.cancel().catch(()=>{});}
    }
    abortCheck(signal);onProgress({phase:'load',loaded:1,total:1});
    this.runtime=new Wllama({default:wasmURL},{suppressNativeLog:true,logger:quiet});
    try {
      await this.runtime.loadModel([file],{n_ctx:2048,n_batch:128,n_ubatch:128,n_threads:Math.min(4,navigator.hardwareConcurrency||2),n_gpu_layers:0,offload_kqv:false,no_kv_offload:true,seed:42});
      abortCheck(signal);this.model=model;
    }catch(e){await this.unload().catch(()=>{});throw e;}
  }
  async generate(question, signal, onText) {
    if(!this.runtime||!this.model)throw new Error('Load a model before starting.');
    let answer='';
    await this.runtime.createChatCompletion({
      messages:recentMessages(this.history,question),stream:true,abortSignal:signal,
      max_tokens:160,temperature:0.6,top_k:40,top_p:0.9,seed:42,cache_prompt:false,
      onData:chunk=>{const text=chunk.choices?.[0]?.delta?.content;if(text){answer+=text;onText(answer);}}
    });
    abortCheck(signal);
    if(!answer.trim())throw new Error('The model did not produce a response. Try a shorter question.');
    this.history=recentMessages(this.history,question).slice(1).concat({role:'assistant',content:answer});
    return answer;
  }
  info(){return {adapter:'wllama 3.13.0',llamaVersion:Wllama.getLibllamaVersion(),model:this.model?.id||null,backend:'WebAssembly CPU',threads:this.runtime?.getNumThreads()||null,crossOriginIsolated:globalThis.crossOriginIsolated===true};}
}