src/tokenizer/cache.ts

import type { FileHash, LoadProgress, RemoteFile } from "./types";

export const TOKENIZER_CACHE = "tokenize-it-assets-v2";
const LEGACY_TOKENIZER_CACHE = "tokenize-it-assets-v1";
const MAX_TOKENIZER_BYTES = 128 * 1024 * 1024;
const CACHE_KEY_PATH = "/.tokenize-it-cache/v2";
let legacyCacheDeletion: Promise<boolean> | null = null;

async function openTokenizerCache(): Promise<Cache> {
  legacyCacheDeletion ??= caches.delete(LEGACY_TOKENIZER_CACHE);
  await legacyCacheDeletion;
  return caches.open(TOKENIZER_CACHE);
}

function cacheKey(hash: FileHash): string {
  return new URL(
    `${CACHE_KEY_PATH}/${hash.algorithm}/${hash.value}`,
    self.location.origin,
  ).href;
}

function hex(bytes: ArrayBuffer): string {
  return Array.from(new Uint8Array(bytes), (byte) =>
    byte.toString(16).padStart(2, "0"),
  ).join("");
}

async function fileDigest(
  bytes: Uint8Array<ArrayBuffer>,
  hash: FileHash,
): Promise<string> {
  if (hash.algorithm === "sha256") {
    return hex(await crypto.subtle.digest("SHA-256", bytes));
  }

  const header = new TextEncoder().encode(`blob ${bytes.byteLength}\0`);
  const gitBlob = new Uint8Array(header.byteLength + bytes.byteLength);
  gitBlob.set(header);
  gitBlob.set(bytes, header.byteLength);
  return hex(await crypto.subtle.digest("SHA-1", gitBlob));
}

async function verifyFileHash(
  bytes: Uint8Array<ArrayBuffer>,
  hash: FileHash,
): Promise<void> {
  if ((await fileDigest(bytes, hash)) !== hash.value) {
    throw new Error("Hugging Face 返回的文件与预期哈希不一致。");
  }
}

async function readResponse(
  response: Response,
  cached: boolean,
  onProgress?: (progress: LoadProgress) => void,
): Promise<Uint8Array<ArrayBuffer>> {
  const contentLength = Number(response.headers.get("Content-Length"));
  const total = Number.isFinite(contentLength) && contentLength > 0 ? contentLength : null;
  if (total !== null && total > MAX_TOKENIZER_BYTES) {
    throw new Error("Tokenizer 文件超过 128 MB 的浏览器加载上限。");
  }

  const reader = response.body?.getReader();
  if (!reader) {
    const bytes = new Uint8Array(await response.arrayBuffer());
    if (bytes.byteLength > MAX_TOKENIZER_BYTES) {
      throw new Error("Tokenizer 文件超过 128 MB 的浏览器加载上限。");
    }
    return bytes;
  }

  const chunks: Uint8Array[] = [];
  let loaded = 0;
  while (true) {
    const { done, value } = await reader.read();
    if (done) break;
    chunks.push(value);
    loaded += value.byteLength;
    if (loaded > MAX_TOKENIZER_BYTES) {
      await reader.cancel();
      throw new Error("Tokenizer 文件超过 128 MB 的浏览器加载上限。");
    }
    onProgress?.({
      phase: "downloading",
      loaded,
      total,
      cached,
      message: cached ? "正在读取浏览器缓存" : "正在下载 Tokenizer",
    });
  }

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

async function fetchFile(
  file: RemoteFile,
  onProgress?: (progress: LoadProgress) => void,
): Promise<Uint8Array<ArrayBuffer>> {
  const cache = await openTokenizerCache();
  const key = cacheKey(file.hash);
  const cachedResponse = await cache.match(key);
  const response = cachedResponse ?? (await fetch(file.url, { mode: "cors" }));

  if (!response.ok) {
    throw new Error(`Tokenizer 文件下载失败（HTTP ${response.status}）。`);
  }

  const bytes = await readResponse(response, Boolean(cachedResponse), onProgress);
  if (cachedResponse) return bytes;

  await verifyFileHash(bytes, file.hash);
  const cacheResponse = new Response(bytes, {
    headers: {
      "Content-Length": String(bytes.byteLength),
      "Content-Type": response.headers.get("Content-Type") ?? "application/octet-stream",
    },
    status: 200,
  });
  // Cache quota exhaustion must not make an otherwise valid tokenizer unusable.
  void cache.put(key, cacheResponse).catch(() => undefined);
  return bytes;
}

export async function fetchTokenizerAsset(
  file: RemoteFile,
  onProgress: (progress: LoadProgress) => void,
): Promise<Uint8Array> {
  return fetchFile(file, onProgress);
}

export async function fetchOptionalJson(
  file: RemoteFile | null,
): Promise<Record<string, unknown> | null> {
  if (!file) return null;
  try {
    const bytes = await fetchFile(file);
    const data = JSON.parse(new TextDecoder().decode(bytes)) as unknown;
    return data && typeof data === "object"
      ? (data as Record<string, unknown>)
      : null;
  } catch {
    return null;
  }
}

async function cachedResponseSize(response: Response): Promise<number> {
  const contentLengthHeader = response.headers.get("Content-Length");
  const contentLength =
    contentLengthHeader === null ? Number.NaN : Number(contentLengthHeader);
  if (Number.isFinite(contentLength) && contentLength >= 0) return contentLength;
  return (await response.blob()).size;
}

export async function tokenizerCacheSize(): Promise<number> {
  const cacheNames = (await caches.keys()).filter(
    (name) => name === TOKENIZER_CACHE || name === LEGACY_TOKENIZER_CACHE,
  );
  const sizes = await Promise.all(
    cacheNames.map(async (name) => {
      const cache = await caches.open(name);
      const requests = await cache.keys();
      const responses = await Promise.all(
        requests.map((request) => cache.match(request)),
      );
      const responseSizes = await Promise.all(
        responses.map((response) =>
          response ? cachedResponseSize(response) : 0,
        ),
      );
      return responseSizes.reduce((total, size) => total + size, 0);
    }),
  );
  return sizes.reduce((total, size) => total + size, 0);
}

export async function clearTokenizerCache(): Promise<void> {
  await Promise.all([
    caches.delete(TOKENIZER_CACHE),
    caches.delete(LEGACY_TOKENIZER_CACHE),
  ]);
}