src/lib/model-cache.js

/**
 * Pulpie WebUI - Model Cache Manager
 *
 * Handles persistent browser caching of sizzlebop/pulpie-orange-small-onnx (model_fp16.onnx)
 * using the CacheStorage API with real-time download progress tracking.
 */

export const LOCAL_MODEL_URL = '/models/sizzlebop/pulpie-orange-small-onnx/model.onnx';
export const REMOTE_MODEL_URL =
  'https://huggingface.co/sizzlebop/pulpie-orange-small-onnx/resolve/main/model.onnx';
export const DEFAULT_MODEL_URL = REMOTE_MODEL_URL;

export const CACHE_NAME = 'pulpie-onnx-v2';
export const MODEL_SIZE_ESTIMATE_BYTES = 847240360; // ~808 MB

/**
 * Determine the best available model URL (local asset vs remote CDN).
 * @returns {Promise<string>}
 */
export async function resolveModelUrl() {
  const isLocal =
    typeof window !== 'undefined' &&
    (window.location.hostname === 'localhost' || window.location.hostname === '127.0.0.1');

  if (isLocal) {
    try {
      const resp = await fetch(LOCAL_MODEL_URL, { method: 'HEAD' });
      const cType = resp.headers.get('content-type') || '';
      const cLen = parseInt(resp.headers.get('content-length') || '0', 10);
      // Ensure it's not an SPA HTML fallback (which is small text/html)
      if (resp.ok && cLen > 10_000_000 && !cType.includes('text/html')) {
        return LOCAL_MODEL_URL;
      }
    } catch {
      // ignore
    }
  }
  return REMOTE_MODEL_URL;
}

/**
 * Check if the model is currently cached in the browser CacheStorage.
 * @param {string} [modelUrl]
 * @returns {Promise<boolean>}
 */
export async function isModelCached(modelUrl = null) {
  if (!('caches' in window)) return false;
  try {
    const cache = await caches.open(CACHE_NAME);
    const targetUrl = modelUrl || (await resolveModelUrl());
    const match = await cache.match(targetUrl);
    return Boolean(match);
  } catch (err) {
    console.warn('[ModelCache] Failed to check cache status:', err);
    return false;
  }
}

/**
 * Retrieve cache metadata (size in MB, cached status).
 * @param {string} [modelUrl]
 * @returns {Promise<{ isCached: boolean, sizeMb: number | null }>}
 */
export async function getCacheInfo(modelUrl = null) {
  if (!('caches' in window)) return { isCached: false, sizeMb: null };
  try {
    const cache = await caches.open(CACHE_NAME);
    const targetUrl = modelUrl || (await resolveModelUrl());
    const match = await cache.match(targetUrl);
    if (!match) return { isCached: false, sizeMb: null };

    const blob = await match.blob();
    const sizeMb = Number((blob.size / (1024 * 1024)).toFixed(1));
    return { isCached: true, sizeMb };
  } catch (err) {
    console.warn('[ModelCache] Failed to get cache info:', err);
    return { isCached: false, sizeMb: null };
  }
}

/**
 * Remove the model from browser CacheStorage.
 * @param {string} [modelUrl]
 * @returns {Promise<boolean>}
 */
export async function clearModelCache(modelUrl = null) {
  if (!('caches' in window)) return false;
  try {
    const cache = await caches.open(CACHE_NAME);
    const targetUrl = modelUrl || (await resolveModelUrl());
    return await cache.delete(targetUrl);
  } catch (err) {
    console.error('[ModelCache] Failed to delete model cache:', err);
    return false;
  }
}

/**
 * Download model with streaming progress reporting and store in CacheStorage.
 *
 * @param {string} [customUrl]
 * @param {function({ loaded: number, total: number, percent: number, speedMbps: number, fromCache: boolean }): void} [onProgress]
 * @returns {Promise<ArrayBuffer>}
 */
export async function loadModelBuffer(customUrl = null, onProgress = null) {
  const modelUrl = customUrl || (await resolveModelUrl());

  let cache = null;
  if ('caches' in window) {
    try {
      cache = await caches.open(CACHE_NAME);
      const cachedResponse = await cache.match(modelUrl);
      if (cachedResponse) {
        if (onProgress) onProgress({ loaded: 1, total: 1, percent: 100, speedMbps: 0, fromCache: true });
        return await cachedResponse.arrayBuffer();
      }
    } catch (cacheErr) {
      console.warn('[ModelCache] Cache read error, falling back to fetch:', cacheErr);
    }
  }

  const startTime = performance.now();
  const response = await fetch(modelUrl, { mode: 'cors' });

  if (!response.ok) {
    throw new Error(`Failed to download model (${response.status} ${response.statusText})`);
  }

  const contentType = response.headers.get('content-type') || '';
  if (contentType.includes('text/html')) {
    throw new Error('Received HTML response instead of model binary file. Please verify model URL.');
  }

  const contentLengthHeader = response.headers.get('content-length');
  const total = contentLengthHeader ? parseInt(contentLengthHeader, 10) : MODEL_SIZE_ESTIMATE_BYTES;

  if (!response.body) {
    const buffer = await response.arrayBuffer();
    if (cache) {
      const responseClone = new Response(buffer.slice(0), {
        headers: { 'Content-Type': 'application/octet-stream', 'Content-Length': String(buffer.byteLength) },
      });
      await cache.put(modelUrl, responseClone).catch(() => {});
    }
    return buffer;
  }

  const reader = response.body.getReader();
  const chunks = [];
  let loaded = 0;
  let lastReportTime = startTime;

  while (true) {
    const { done, value } = await reader.read();
    if (done) break;

    chunks.push(value);
    loaded += value.length;

    const now = performance.now();
    if (onProgress && now - lastReportTime > 80) {
      const elapsedSec = (now - startTime) / 1000;
      const speedMbps = elapsedSec > 0 ? Number(((loaded * 8) / (elapsedSec * 1024 * 1024)).toFixed(2)) : 0;
      const percent = Math.min(100, Number(((loaded / total) * 100).toFixed(1)));
      onProgress({ loaded, total, percent, speedMbps, fromCache: false });
      lastReportTime = now;
    }
  }

  const modelBuffer = new Uint8Array(loaded);
  let offset = 0;
  for (const chunk of chunks) {
    modelBuffer.set(chunk, offset);
    offset += chunk.length;
  }

  if (onProgress) {
    onProgress({ loaded, total: loaded, percent: 100, speedMbps: 0, fromCache: false });
  }

  if (cache) {
    try {
      const cacheResponse = new Response(modelBuffer.buffer.slice(0), {
        headers: {
          'Content-Type': 'application/octet-stream',
          'Content-Length': String(loaded),
        },
      });
      await cache.put(modelUrl, cacheResponse);
    } catch (putErr) {
      console.warn('[ModelCache] Could not persist model to CacheStorage:', putErr);
    }
  }

  return modelBuffer.buffer;
}