import { stft, type Stft } from './fft';
import { css, drawAxes, fmtNum, PAD, setupCanvas } from './plot';

function hexToRgb(h: string): [number, number, number] {
  const v = parseInt(h.replace('#', ''), 16);
  return [(v >> 16) & 255, (v >> 8) & 255, v & 255];
}

// Sequential ramp read from a CSS custom property: comma-separated hex steps.
export function rampLUT(varName: string): Uint8ClampedArray {
  const stops = css(varName).split(',').map((s) => hexToRgb(s.trim()));
  const lut = new Uint8ClampedArray(256 * 3);
  for (let i = 0; i < 256; i++) {
    const t = (i / 255) * (stops.length - 1);
    const k = Math.min(stops.length - 2, Math.floor(t));
    const f = t - k;
    for (let c = 0; c < 3; c++) lut[3 * i + c] = stops[k][c] + f * (stops[k + 1][c] - stops[k][c]);
  }
  return lut;
}

export class Spectrogram {
  private s: Stft | null = null;
  fmax = 5000;
  range = 80;
  private tip: HTMLDivElement;

  constructor(private canvas: HTMLCanvasElement) {
    this.tip = document.createElement('div');
    this.tip.className = 'tooltip';
    this.tip.hidden = true;
    canvas.parentElement!.append(this.tip);
    canvas.addEventListener('mousemove', (e) => this.hover(e.offsetX, e.offsetY));
    canvas.addEventListener('mouseleave', () => (this.tip.hidden = true));
    new ResizeObserver(() => this.draw()).observe(canvas);
  }

  setSignal(x: Float32Array, fs: number) {
    this.s = stft(x, fs, 2048, Math.max(64, Math.floor(x.length / 1200)));
    this.draw();
  }

  draw() {
    const { ctx, w, h } = setupCanvas(this.canvas);
    ctx.clearRect(0, 0, w, h);
    const s = this.s;
    if (!s) return;
    const fmax = Math.min(this.fmax, s.fs / 2);
    const nb = Math.max(2, Math.round((fmax / s.fs) * s.nfft));
    const nf = s.mags.length;
    let top = -Infinity;
    for (const m of s.mags) for (let k = 0; k <= nb; k++) top = Math.max(top, m[k]);
    const lut = rampLUT('--ramp');
    const img = new ImageData(nf, nb);
    for (let j = 0; j < nf; j++) {
      const m = s.mags[j];
      for (let k = 0; k < nb; k++) {
        const v = Math.max(0, Math.min(1, (m[k] - top + this.range) / this.range));
        const li = Math.round(v * 255) * 3;
        const p = 4 * ((nb - 1 - k) * nf + j);
        img.data[p] = lut[li];
        img.data[p + 1] = lut[li + 1];
        img.data[p + 2] = lut[li + 2];
        img.data[p + 3] = 255;
      }
    }
    const off = document.createElement('canvas');
    off.width = nf;
    off.height = nb;
    off.getContext('2d')!.putImageData(img, 0, 0);
    const duration = ((nf - 1) * s.hop) / s.fs;
    ctx.imageSmoothingEnabled = true;
    ctx.drawImage(off, PAD.l, PAD.t, w - PAD.l - PAD.r, h - PAD.t - PAD.b);
    drawAxes(ctx, w, h, [0, duration], [0, fmax], { xLabel: 'time (s)', yLabel: 'frequency (Hz)', grid: false });
  }

  private hover(px: number, py: number) {
    const s = this.s;
    if (!s) return;
    const w = this.canvas.clientWidth;
    const h = this.canvas.clientHeight;
    if (px < PAD.l || px > w - PAD.r || py < PAD.t || py > h - PAD.b) {
      this.tip.hidden = true;
      return;
    }
    const fmax = Math.min(this.fmax, s.fs / 2);
    const fx = (px - PAD.l) / (w - PAD.l - PAD.r);
    const fy = 1 - (py - PAD.t) / (h - PAD.t - PAD.b);
    const j = Math.min(s.mags.length - 1, Math.round(fx * (s.mags.length - 1)));
    const f = fy * fmax;
    const k = Math.round((f / s.fs) * s.nfft);
    this.tip.textContent = `t: ${fmtNum((j * s.hop) / s.fs)} s\nf: ${fmtNum(f)} Hz\nlevel: ${s.mags[j][k].toFixed(1)} dB`;
    this.tip.hidden = false;
    this.tip.style.left = `${Math.min(px + 12, w - 140)}px`;
    this.tip.style.top = `${PAD.t + 4}px`;
  }
}