import { createSHA256 } from 'hash-wasm';
import manifest from './model-manifest.json';

export const REVISION = '04a6c49fcba3a65a0351c92644c3a7e9d4343059';
export const MODEL_ID = 'Mike0021/MiniCPM5-2B-ONNX';
export const MODEL_BYTES = Object.values(manifest.files).reduce((n, f) => n + f.bytes, 0);
const CHUNK = 32 * 1024 ** 2;
const sleep = (ms, signal) => new Promise((resolve, reject) => {
  signal?.throwIfAborted();
  const abort = () => { clearTimeout(timer); reject(signal.reason); };
  const timer = setTimeout(() => { signal?.removeEventListener('abort', abort); resolve(); }, ms);
  signal?.addEventListener('abort', abort, { once: true });
});

async function fetchPart(url, start, end, total, signal) {
  for (let attempt = 0; ; attempt++) {
    try {
      const part = total > CHUNK;
      const response = await fetch(url, { headers: part ? { Range: `bytes=${start}-${end}` } : {},
        signal: AbortSignal.any([signal, AbortSignal.timeout(90_000)]), cache: 'no-store' });
      if (!response.ok) throw Error(`Download returned HTTP ${response.status}`);
      if (part && response.status !== 206) {
        await response.body?.cancel();
        throw Error('The model host did not honor byte ranges.');
      }
      const range = response.headers.get('Content-Range');
      if (range && range !== `bytes ${start}-${end}/${total}`) {
        await response.body?.cancel(); throw Error('Incorrect download range.');
      }
      const data = new Uint8Array(await response.arrayBuffer());
      if (data.length !== end - start + 1) throw Error('Incomplete download chunk.');
      return data;
    } catch (error) {
      signal.throwIfAborted();
      if (attempt === 2) throw error;
      await sleep(500 * 2 ** attempt, signal);
    }
  }
}

export async function prepareModelCache(baseURL, options) {
  if (navigator.locks) return navigator.locks.request('pi-minicpm5-download-' + REVISION,
    { signal: options.signal }, () => downloadModel(baseURL, options));
  return downloadModel(baseURL, options);
}

async function readCacheState(create = false) {
  if (!navigator.storage?.getDirectory) throw Error('This app needs browser file storage (OPFS) to cache the 1.84 GB model. Open it in a regular browser tab.');
  const root = await (await navigator.storage.getDirectory()).getDirectoryHandle('pi-minicpm5-' + REVISION, { create });
  const receiptsFile = await root.getFileHandle('verified.json', { create });
  let receipts;
  try { receipts = JSON.parse(await (await receiptsFile.getFile()).text()); } catch { receipts = {}; }
  if (!receipts || typeof receipts !== 'object' || Array.isArray(receipts)) receipts = {};
  let cached = 0;
  const pending = [];
  const key = file => encodeURIComponent(file);
  for (const [file, spec] of Object.entries(manifest.files)) {
    try {
      const stored = await (await root.getFileHandle(key(file))).getFile();
      const receipt = receipts[file];
      if (stored.size === spec.bytes && receipt?.sha256 === spec.sha256 && receipt.modified === stored.lastModified) {
        cached += spec.bytes; continue;
      }
    } catch (e) { if (e.name !== 'NotFoundError') throw e; }
    pending.push([file, spec]);
  }
  return { root, receiptsFile, receipts, pending, cached };
}

export async function isModelCached() {
  try { return (await readCacheState()).pending.length === 0; }
  catch { return false; }
}

function cacheMiss() {
  return Object.assign(Error('The saved model is incomplete. Confirm the download to restore it.'), { code: 'MODEL_CACHE_MISS' });
}

async function downloadModel(baseURL, { signal, cachedOnly = false, onProgress = () => {} }) {
  let state;
  try { state = await readCacheState(!cachedOnly); }
  catch (error) { if (cachedOnly && error.name === 'NotFoundError') throw cacheMiss(); throw error; }
  const { root, receiptsFile, receipts, pending, cached } = state;
  // Recheck under the download lock: another tab or eviction may have changed
  // the cache after the UI checked it. Automatic restores never fetch weights.
  if (cachedOnly && pending.length) throw cacheMiss();
  let completed = cached;
  const key = file => encodeURIComponent(file);
  onProgress({ phase: 'download', loaded: completed, total: MODEL_BYTES, cached });
  if (pending.length) {
    const estimate = await navigator.storage.estimate();
    if (estimate.quota && estimate.quota - estimate.usage < MODEL_BYTES - completed + CHUNK) {
      throw Error('Not enough browser storage for the model. Free at least 2 GB and retry.');
    }
  }
  const controller = new AbortController();
  const combined = AbortSignal.any([signal, controller.signal]);
  let receiptWrites = Promise.resolve();
  const saveReceipts = () => {
    const serialized = JSON.stringify(receipts);
    const write = receiptWrites.then(async () => {
      const stream = await receiptsFile.createWritable(); await stream.write(serialized); await stream.close();
    });
    receiptWrites = write.catch(() => {}); return write;
  };
  async function worker() {
    while (pending.length) {
      combined.throwIfAborted();
      const [file, spec] = pending.shift();
      const handle = await root.getFileHandle(key(file), { create: true });
      const output = await handle.createWritable();
      try {
        const hash = await createSHA256(); hash.init();
        for (let start = 0; start < spec.bytes; start += CHUNK) {
          const end = Math.min(start + CHUNK, spec.bytes) - 1;
          const data = await fetchPart(baseURL + file, start, end, spec.bytes, combined);
          combined.throwIfAborted();
          hash.update(data); await output.write(data); completed += data.length;
          onProgress({ phase: 'download', file, loaded: completed, total: MODEL_BYTES, cached });
        }
        if (hash.digest('hex') !== spec.sha256) throw Error('SHA-256 verification failed: ' + file);
        await output.close();
        receipts[file] = { sha256: spec.sha256, modified: (await handle.getFile()).lastModified };
        await saveReceipts();
      } catch (error) { await output.abort().catch(() => {}); controller.abort(error); throw error; }
    }
  }
  const outcomes = await Promise.allSettled(Array.from({ length: 6 }, worker));
  const failed = outcomes.find(r => r.status === 'rejected');
  if (failed) throw failed.reason;
  await receiptWrites;
  return {
    cachedBytes: cached,
    async match(request) {
      const url = typeof request === 'string' ? request : request.url;
      if (!url.startsWith(baseURL)) return undefined;
      const file = url.slice(baseURL.length);
      if (!Object.hasOwn(manifest.files, file)) return undefined;
      const stored = await (await root.getFileHandle(key(file))).getFile();
      return new Response(stored, { headers: { 'Content-Length': String(stored.size) } });
    },
    async put() { /* All supported files are already downloaded and verified. */ },
  };
}