wrongmove/frontend/src/workers/HexgridHeatmapClient.ts

import RBush from 'rbush';
import type {
  FeatureItem,
  WorkerResponse,
  ColorScaleResultMessage,
  BoundsResultMessage,
  GridResultMessage,
} from './types';
import { toFeatureItem } from './types';
import { record } from '../services/perfCollector';

interface HexgridConfig {
  intensity: number;
  spread: number;
  cellDensity: number;
  propertyName: string;
}

interface PendingPromise<T> {
  resolve: (value: T) => void;
  reject: (reason: unknown) => void;
}

export class HexgridHeatmapClient {
  private map: mapboxgl.Map;
  private worker: Worker;
  private mainTree = new RBush<FeatureItem>();
  private requestCounter = 0;

  // Layer swap state (matches original HexgridHeatmap)
  private layerA: string;
  private layerB: string;
  private sourceA: mapboxgl.GeoJSONSource | null = null;
  private sourceB: mapboxgl.GeoJSONSource | null = null;
  private activeIsA = true;
  private lastZoom: number | null = null;
  private clearTimeout: ReturnType<typeof setTimeout> | null = null;

  // Grid calculation coalescing
  private calculatingGrid = false;
  private recalcWhenReady = false;

  // Pending grid result tracking
  private pendingGridRequestId: number | null = null;

  // Pending promises for async operations
  private pendingColorScale = new Map<number, PendingPromise<ColorScaleResultMessage>>();
  private pendingBounds = new Map<number, PendingPromise<BoundsResultMessage>>();

  // Send timestamps for round-trip measurement
  private sendTimestamps = new Map<number, number>();

  private config: HexgridConfig = {
    intensity: 8,
    spread: 0.1,
    cellDensity: 1,
    propertyName: 'count',
  };

  private moveEndHandler: () => void;

  constructor(map: mapboxgl.Map, layername: string, addBefore?: string) {
    this.map = map;
    this.layerA = layername;
    this.layerB = layername + '-back';

    this.setupLayers(layername, addBefore);

    // Create worker
    this.worker = new Worker(
      new URL('./hexgrid.worker.ts', import.meta.url),
      { type: 'module' }
    );
    this.worker.onmessage = (e: MessageEvent<WorkerResponse>) => {
      this.handleWorkerMessage(e.data);
    };

    // Bind moveend
    this.moveEndHandler = () => this.updateGrid();
    this.map.on('moveend', this.moveEndHandler);
  }

  private setupLayers(layername: string, addBefore?: string): void {
    const defaultPaint: mapboxgl.FillPaint = {
      'fill-opacity': 0,
      'fill-opacity-transition': { duration: 200, delay: 0 },
      'fill-color': {
        property: 'count',
        stops: [
          [0, 'rgba(0,185,243,0)'],
          [50, 'rgba(0,185,243,0.24)'],
          [130, 'rgba(255,223,0,0.3)'],
          [200, 'rgba(255,105,0,0.3)'],
        ],
      },
    };
    const emptyData: GeoJSON.FeatureCollection = { type: 'FeatureCollection', features: [] };

    // Layer A: primary, starts active at opacity 1
    this.map.addLayer({
      id: layername,
      type: 'fill',
      source: { type: 'geojson', data: emptyData },
      paint: { ...defaultPaint, 'fill-opacity': 1.0 },
    }, addBefore);

    // Layer B: back layer, starts inactive at opacity 0, renders below A
    this.map.addLayer({
      id: layername + '-back',
      type: 'fill',
      source: { type: 'geojson', data: emptyData },
      paint: { ...defaultPaint },
    }, layername);

    this.sourceA = this.map.getSource(layername) as mapboxgl.GeoJSONSource;
    this.sourceB = this.map.getSource(layername + '-back') as mapboxgl.GeoJSONSource;
  }

  private nextRequestId(): number {
    return ++this.requestCounter;
  }

  private handleWorkerMessage(msg: WorkerResponse): void {
    switch (msg.type) {
      case 'DATA_READY': {
        const sent = this.sendTimestamps.get(msg.requestId);
        if (sent !== undefined) {
          this.sendTimestamps.delete(msg.requestId);
          record('worker_roundtrip', 'set_data', (performance.now() - sent) / 1000);
        }
        record('worker_compute', 'set_data', msg.computeMs / 1000);
        break;
      }

      case 'GRID_RESULT':
        this.handleGridResult(msg);
        break;

      case 'COLOR_SCALE_RESULT': {
        const pending = this.pendingColorScale.get(msg.requestId);
        if (pending) {
          this.pendingColorScale.delete(msg.requestId);
          pending.resolve(msg);
        }
        break;
      }

      case 'BOUNDS_RESULT': {
        const pending = this.pendingBounds.get(msg.requestId);
        if (pending) {
          this.pendingBounds.delete(msg.requestId);
          pending.resolve(msg);
        }
        break;
      }
    }
  }

  private handleGridResult(msg: GridResultMessage): void {
    // Discard stale grid results
    if (msg.requestId !== this.pendingGridRequestId) return;
    this.pendingGridRequestId = null;

    // Record perf metrics
    const sent = this.sendTimestamps.get(msg.requestId);
    if (sent !== undefined) {
      this.sendTimestamps.delete(msg.requestId);
      record('worker_roundtrip', 'generate_grid', (performance.now() - sent) / 1000);
    }
    record('worker_compute', 'generate_grid', msg.computeMs / 1000);

    const hexgrid = msg.hexgrid as unknown as GeoJSON.FeatureCollection;
    const currentZoom = Math.floor(this.map.getZoom());
    const zoomChanged = this.lastZoom !== null && this.lastZoom !== currentZoom;
    this.lastZoom = currentZoom;

    const activeSource = this.activeIsA ? this.sourceA! : this.sourceB!;
    const activeLayer = this.activeIsA ? this.layerA : this.layerB;
    const inactiveSource = this.activeIsA ? this.sourceB! : this.sourceA!;
    const inactiveLayer = this.activeIsA ? this.layerB : this.layerA;

    // Cancel any pending cleanup timeout
    if (this.clearTimeout) {
      clearTimeout(this.clearTimeout);
      this.clearTimeout = null;
    }

    if (zoomChanged) {
      // Cross-fade between layers
      this.map.moveLayer(activeLayer, inactiveLayer);
      inactiveSource.setData(hexgrid);
      this.map.setPaintProperty(inactiveLayer, 'fill-opacity', 1.0);
      this.map.setPaintProperty(activeLayer, 'fill-opacity', 0);
      this.activeIsA = !this.activeIsA;

      // Clear old layer data after transition
      const src = activeSource;
      this.clearTimeout = setTimeout(() => {
        this.clearTimeout = null;
        src.setData({ type: 'FeatureCollection', features: [] });
      }, 250);
    } else {
      // Pan only: swap data directly on active source
      activeSource.setData(hexgrid);
    }

    this.calculatingGrid = false;
    if (this.recalcWhenReady) {
      this.recalcWhenReady = false;
      this.updateGrid();
    }
  }

  // --- Public API ---

  // eslint-disable-next-line @typescript-eslint/no-explicit-any
  setData(data: { features: Array<{ properties: any; geometry: { type: 'Point'; coordinates: [number, number] } }> }): void {
    // Build main-thread R-tree for click queries
    const t0 = performance.now();
    const items = data.features.map(toFeatureItem);
    this.mainTree.clear();
    this.mainTree.load(items);
    record('main_thread', 'rtree_build', (performance.now() - t0) / 1000);
    record('feature_count', 'set_data', data.features.length);

    // Post to worker for grid generation tree
    const requestId = this.nextRequestId();
    this.sendTimestamps.set(requestId, performance.now());
    this.worker.postMessage({
      type: 'SET_DATA',
      requestId,
      features: data.features,
    });
  }

  setPropertyName(name: string): void {
    this.config.propertyName = name;
  }

  setSpread(v: number): void {
    this.config.spread = v;
  }

  setCellDensity(v: number): void {
    this.config.cellDensity = v;
  }

  setIntensity(v: number): void {
    this.config.intensity = v;
  }

  setColorStops(stops: [number, string][]): void {
    const colorProp = { property: 'count', stops };
    this.map.setPaintProperty(this.layerA, 'fill-color', colorProp);
    this.map.setPaintProperty(this.layerB, 'fill-color', colorProp);
  }

  update(): void {
    this.updateGrid();
  }

  searchTree(bounds: { minX: number; minY: number; maxX: number; maxY: number }): FeatureItem[] {
    return this.mainTree.search(bounds);
  }

  computeColorScale(
    metricMode: string,
    percentileConfig: { minBound: number; maxBound: number }
  ): Promise<ColorScaleResultMessage> {
    const requestId = this.nextRequestId();
    this.sendTimestamps.set(requestId, performance.now());
    return new Promise((resolve, reject) => {
      this.pendingColorScale.set(requestId, {
        resolve: (msg) => {
          const sent = this.sendTimestamps.get(requestId);
          if (sent !== undefined) {
            this.sendTimestamps.delete(requestId);
            record('worker_roundtrip', 'color_scale', (performance.now() - sent) / 1000);
          }
          record('worker_compute', 'color_scale', msg.computeMs / 1000);
          resolve(msg);
        },
        reject,
      });
      this.worker.postMessage({
        type: 'COMPUTE_COLOR_SCALE',
        requestId,
        metricMode,
        percentileConfig,
      });
    });
  }

  computeBounds(
    boundsPercentileConfig: { clipMin: number; clipMax: number }
  ): Promise<BoundsResultMessage> {
    const requestId = this.nextRequestId();
    this.sendTimestamps.set(requestId, performance.now());
    return new Promise((resolve, reject) => {
      this.pendingBounds.set(requestId, {
        resolve: (msg) => {
          const sent = this.sendTimestamps.get(requestId);
          if (sent !== undefined) {
            this.sendTimestamps.delete(requestId);
            record('worker_roundtrip', 'bounds', (performance.now() - sent) / 1000);
          }
          record('worker_compute', 'bounds', msg.computeMs / 1000);
          resolve(msg);
        },
        reject,
      });
      this.worker.postMessage({
        type: 'COMPUTE_BOUNDS',
        requestId,
        boundsPercentileConfig,
      });
    });
  }

  destroy(): void {
    this.map.off('moveend', this.moveEndHandler);
    if (this.clearTimeout) {
      clearTimeout(this.clearTimeout);
    }
    // Reject any pending promises
    for (const [, p] of this.pendingColorScale) {
      p.reject(new Error('HexgridHeatmapClient destroyed'));
    }
    for (const [, p] of this.pendingBounds) {
      p.reject(new Error('HexgridHeatmapClient destroyed'));
    }
    this.pendingColorScale.clear();
    this.pendingBounds.clear();
    this.worker.terminate();
  }

  // --- Private ---

  private updateGrid(): void {
    if (!this.calculatingGrid) {
      this.calculatingGrid = true;

      const zoom = this.map.getZoom();
      const b = this.map.getBounds()!;
      const bounds: [number, number, number, number] = [
        b.getWest(), b.getSouth(), b.getEast(), b.getNorth(),
      ];

      const requestId = this.nextRequestId();
      this.pendingGridRequestId = requestId;
      this.sendTimestamps.set(requestId, performance.now());

      this.worker.postMessage({
        type: 'GENERATE_GRID',
        requestId,
        zoom,
        bounds,
        config: this.config,
      });
    } else {
      this.recalcWhenReady = true;
    }
  }
}