next-tjwater-frontend/src/features/workbench/map/city-pmtiles-cache.ts

import {
  EtagMismatch,
  FetchSource,
  PMTiles,
  type Protocol,
  type RangeResponse,
  type Source
} from "pmtiles";

const CACHE_DATABASE_NAME = "tjwater-city-pmtiles-cache";
const CACHE_DATABASE_VERSION = 2;
const CACHE_STORE_NAME = "ranges";
const CACHE_URL_INDEX = "url";
const CACHE_ARCHIVE_STORE_NAME = "archives";
const MAX_ARCHIVE_CACHE_BYTES = 48 * 1024 * 1024;
const PMTILES_HEADER_SIZE = 127;

export interface CityPmtilesRangeRecord extends RangeResponse {
  key: string;
  url: string;
  revision: string;
  offset: number;
  length: number;
  storedAt: number;
}

export interface CityPmtilesRangeStore {
  get(key: string): Promise<CityPmtilesRangeRecord | undefined>;
  put(record: CityPmtilesRangeRecord): Promise<void>;
  prepareRevision(url: string, revision: string): Promise<void>;
  deleteRevision(url: string, revision: string): Promise<void>;
}

interface CityPmtilesArchiveRecord {
  url: string;
  revision: string;
  cachedBytes: number;
}

function abortIfRequested(signal?: AbortSignal) {
  if (signal?.aborted) {
    throw signal.reason ?? new DOMException("The operation was aborted", "AbortError");
  }
}

function requestResult<T>(request: IDBRequest<T>): Promise<T> {
  return new Promise((resolve, reject) => {
    request.addEventListener("success", () => resolve(request.result), { once: true });
    request.addEventListener("error", () => reject(request.error), { once: true });
  });
}

function transactionDone(transaction: IDBTransaction): Promise<void> {
  return new Promise((resolve, reject) => {
    transaction.addEventListener("complete", () => resolve(), { once: true });
    transaction.addEventListener(
      "abort",
      () => reject(transaction.error ?? new Error("City PMTiles cache transaction aborted")),
      { once: true }
    );
    transaction.addEventListener(
      "error",
      () => reject(transaction.error ?? new Error("City PMTiles cache transaction failed")),
      { once: true }
    );
  });
}

class IndexedDbCityPmtilesRangeStore implements CityPmtilesRangeStore {
  private readonly databasePromise: Promise<IDBDatabase>;

  constructor(indexedDb: IDBFactory) {
    this.databasePromise = new Promise((resolve, reject) => {
      const request = indexedDb.open(CACHE_DATABASE_NAME, CACHE_DATABASE_VERSION);
      request.addEventListener(
        "upgradeneeded",
        () => {
          const database = request.result;
          const transaction = request.transaction!;
          const store = database.objectStoreNames.contains(CACHE_STORE_NAME)
            ? transaction.objectStore(CACHE_STORE_NAME)
            : database.createObjectStore(CACHE_STORE_NAME, { keyPath: "key" });
          if (!store.indexNames.contains(CACHE_URL_INDEX)) {
            store.createIndex(CACHE_URL_INDEX, "url", { unique: false });
          }
          if (!database.objectStoreNames.contains(CACHE_ARCHIVE_STORE_NAME)) {
            database.createObjectStore(CACHE_ARCHIVE_STORE_NAME, { keyPath: "url" });
          }
        },
        { once: true }
      );
      request.addEventListener(
        "success",
        () => {
          const database = request.result;
          database.addEventListener("versionchange", () => database.close());
          resolve(database);
        },
        { once: true }
      );
      request.addEventListener("error", () => reject(request.error), { once: true });
      request.addEventListener(
        "blocked",
        () => reject(new Error("City PMTiles cache database upgrade was blocked")),
        { once: true }
      );
    });
  }

  async get(key: string) {
    const database = await this.databasePromise;
    const transaction = database.transaction(CACHE_STORE_NAME, "readonly");
    const result = requestResult(
      transaction.objectStore(CACHE_STORE_NAME).get(key) as IDBRequest<
        CityPmtilesRangeRecord | undefined
      >
    );
    const [record] = await Promise.all([result, transactionDone(transaction)]);
    return record;
  }

  async put(record: CityPmtilesRangeRecord) {
    const database = await this.databasePromise;
    const transaction = database.transaction(
      [CACHE_STORE_NAME, CACHE_ARCHIVE_STORE_NAME],
      "readwrite"
    );
    const ranges = transaction.objectStore(CACHE_STORE_NAME);
    const archives = transaction.objectStore(CACHE_ARCHIVE_STORE_NAME);
    const archiveRequest = archives.get(record.url) as IDBRequest<
      CityPmtilesArchiveRecord | undefined
    >;
    archiveRequest.addEventListener("success", () => {
      const archive = archiveRequest.result;
      if (!archive || archive.revision !== record.revision) return;
      const existingRequest = ranges.get(record.key) as IDBRequest<
        CityPmtilesRangeRecord | undefined
      >;
      existingRequest.addEventListener("success", () => {
        const cachedBytes =
          archive.cachedBytes -
          (existingRequest.result?.data.byteLength ?? 0) +
          record.data.byteLength;
        if (cachedBytes > MAX_ARCHIVE_CACHE_BYTES) return;
        ranges.put(record);
        archives.put({ ...archive, cachedBytes });
      });
    });
    await transactionDone(transaction);
  }

  async prepareRevision(url: string, revision: string) {
    const database = await this.databasePromise;
    const transaction = database.transaction(
      [CACHE_STORE_NAME, CACHE_ARCHIVE_STORE_NAME],
      "readwrite"
    );
    const ranges = transaction.objectStore(CACHE_STORE_NAME);
    const archives = transaction.objectStore(CACHE_ARCHIVE_STORE_NAME);
    const archiveRequest = archives.get(url) as IDBRequest<CityPmtilesArchiveRecord | undefined>;
    archiveRequest.addEventListener("success", () => {
      if (archiveRequest.result?.revision === revision) return;
      let cachedBytes = 0;
      const cursorRequest = ranges.index(CACHE_URL_INDEX).openCursor(IDBKeyRange.only(url));
      cursorRequest.addEventListener("success", () => {
        const cursor = cursorRequest.result;
        if (!cursor) {
          archives.put({ url, revision, cachedBytes });
          return;
        }
        const record = cursor.value as CityPmtilesRangeRecord;
        if (record.revision === revision) cachedBytes += record.data.byteLength;
        else cursor.delete();
        cursor.continue();
      });
    });
    await transactionDone(transaction);
  }

  async deleteRevision(url: string, revision: string) {
    const database = await this.databasePromise;
    const transaction = database.transaction(
      [CACHE_STORE_NAME, CACHE_ARCHIVE_STORE_NAME],
      "readwrite"
    );
    const ranges = transaction.objectStore(CACHE_STORE_NAME);
    const archives = transaction.objectStore(CACHE_ARCHIVE_STORE_NAME);
    const index = ranges.index(CACHE_URL_INDEX);
    const request = index.openCursor(IDBKeyRange.only(url));
    request.addEventListener("success", () => {
      const cursor = request.result;
      if (!cursor) {
        const archiveRequest = archives.get(url) as IDBRequest<
          CityPmtilesArchiveRecord | undefined
        >;
        archiveRequest.addEventListener("success", () => {
          if (archiveRequest.result?.revision === revision) {
            archives.put({ url, revision, cachedBytes: 0 });
          }
        });
        return;
      }
      const record = cursor.value as CityPmtilesRangeRecord;
      if (record.revision === revision) cursor.delete();
      cursor.continue();
    });
    await transactionDone(transaction);
  }
}

let sharedPersistentStore: CityPmtilesRangeStore | null | undefined;

function createPersistentStore(): CityPmtilesRangeStore | undefined {
  if (sharedPersistentStore === undefined) {
    sharedPersistentStore =
      "indexedDB" in globalThis && globalThis.indexedDB
        ? new IndexedDbCityPmtilesRangeStore(globalThis.indexedDB)
        : null;
  }
  return sharedPersistentStore ?? undefined;
}

function rangeKey(url: string, revision: string, offset: number, length: number) {
  return JSON.stringify([url, revision, offset, length]);
}

function isValidRangeLength(offset: number, requestedLength: number, actualLength: number) {
  return (
    actualLength === requestedLength ||
    (offset === 0 && actualLength >= PMTILES_HEADER_SIZE && actualLength < requestedLength)
  );
}

/**
 * Persists the exact byte ranges PMTiles requests. The manifest revision keeps
 * an updated archive from reusing stale header, directory, or tile bytes.
 */
export class PersistentCityPmtilesSource implements Source {
  private store: CityPmtilesRangeStore | undefined;
  private readonly prepareStore: Promise<void> | undefined;

  constructor(
    private readonly url: string,
    private readonly revision: string,
    private readonly upstream: Source = new FetchSource(url),
    store: CityPmtilesRangeStore | null | undefined = createPersistentStore()
  ) {
    this.store = store ?? undefined;
    this.prepareStore = this.store
      ?.prepareRevision(this.url, this.revision)
      .catch(() => this.disableStore());
  }

  getKey() {
    return this.url;
  }

  async getBytes(offset: number, length: number, signal?: AbortSignal, etag?: string) {
    abortIfRequested(signal);
    await this.prepareStore;
    abortIfRequested(signal);

    const key = rangeKey(this.url, this.revision, offset, length);
    const cached = await this.readCachedRange(key);
    abortIfRequested(signal);
    if (
      cached &&
      isValidRangeLength(offset, length, cached.data.byteLength) &&
      (!etag || cached.etag === etag)
    ) {
      return {
        data: cached.data.slice(0),
        etag: cached.etag,
        cacheControl: cached.cacheControl,
        expires: cached.expires
      };
    }

    let response: RangeResponse;
    try {
      response = await this.upstream.getBytes(offset, length, signal, etag);
    } catch (error) {
      if (error instanceof EtagMismatch) await this.clearCurrentRevision();
      throw error;
    }
    abortIfRequested(signal);
    if (!isValidRangeLength(offset, length, response.data.byteLength)) {
      throw new Error(
        `City PMTiles range ${offset}-${offset + length - 1} returned ${response.data.byteLength} bytes instead of ${length}`
      );
    }

    void this.writeCachedRange({
      ...response,
      data: response.data.slice(0),
      etag: response.etag ?? etag,
      key,
      url: this.url,
      revision: this.revision,
      offset,
      length,
      storedAt: Date.now()
    });
    return response;
  }

  private async readCachedRange(key: string) {
    if (!this.store) return undefined;
    try {
      return await this.store.get(key);
    } catch {
      this.disableStore();
      return undefined;
    }
  }

  private async writeCachedRange(record: CityPmtilesRangeRecord) {
    if (!this.store) return;
    try {
      await this.store.put(record);
    } catch (error) {
      if (error instanceof DOMException && error.name === "QuotaExceededError") {
        await this.clearCurrentRevision();
      }
      this.disableStore();
    }
  }

  private async clearCurrentRevision() {
    if (!this.store) return;
    try {
      await this.store.deleteRevision(this.url, this.revision);
    } catch {
      this.disableStore();
    }
  }

  private disableStore() {
    this.store = undefined;
  }
}

export function registerCityPmtiles(protocol: Protocol, url: string, revision?: string) {
  protocol.add(
    new PMTiles(revision ? new PersistentCityPmtilesSource(url, revision) : new FetchSource(url))
  );
}