audio-ui/components/analysis/Spectrogram.tsx

'use client';

import * as React from 'react';
import { cn } from '@/lib/utils/cn';

export interface SpectrogramProps {
  analyserNode: AnalyserNode | null;
  className?: string;
}

export function Spectrogram({ analyserNode, className }: SpectrogramProps) {
  const canvasRef = React.useRef<HTMLCanvasElement>(null);
  const animationFrameRef = React.useRef<number | undefined>(undefined);
  const spectrogramDataRef = React.useRef<ImageData | null>(null);

  React.useEffect(() => {
    if (!analyserNode || !canvasRef.current) return;

    const canvas = canvasRef.current;
    const ctx = canvas.getContext('2d');
    if (!ctx) return;

    // Set canvas size
    const dpr = window.devicePixelRatio || 1;
    const rect = canvas.getBoundingClientRect();
    canvas.width = rect.width * dpr;
    canvas.height = rect.height * dpr;
    ctx.scale(dpr, dpr);

    const bufferLength = analyserNode.frequencyBinCount;
    const dataArray = new Uint8Array(bufferLength);

    // Initialize spectrogram data
    spectrogramDataRef.current = ctx.createImageData(rect.width, rect.height);

    const draw = () => {
      animationFrameRef.current = requestAnimationFrame(draw);

      analyserNode.getByteFrequencyData(dataArray);

      if (!spectrogramDataRef.current) return;

      const imageData = spectrogramDataRef.current;

      // Shift existing data to the left by 1 pixel
      for (let y = 0; y < rect.height; y++) {
        for (let x = 0; x < rect.width - 1; x++) {
          const sourceIndex = ((y * rect.width) + x + 1) * 4;
          const targetIndex = ((y * rect.width) + x) * 4;
          imageData.data[targetIndex] = imageData.data[sourceIndex];
          imageData.data[targetIndex + 1] = imageData.data[sourceIndex + 1];
          imageData.data[targetIndex + 2] = imageData.data[sourceIndex + 2];
          imageData.data[targetIndex + 3] = imageData.data[sourceIndex + 3];
        }
      }

      // Add new column on the right
      const x = rect.width - 1;
      for (let y = 0; y < rect.height; y++) {
        // Map frequency bins to canvas height (inverted)
        const freqIndex = Math.floor((1 - y / rect.height) * bufferLength);
        const value = dataArray[freqIndex];

        // Color mapping with transparency: transparent (0) -> blue -> cyan -> green -> yellow -> red (255)
        let r, g, b, a;
        if (value < 64) {
          // Transparent to blue
          const t = value / 64;
          r = 0;
          g = 0;
          b = Math.round(255 * t);
          a = Math.round(255 * t);
        } else if (value < 128) {
          // Blue to cyan
          r = 0;
          g = (value - 64) * 4;
          b = 255;
          a = 255;
        } else if (value < 192) {
          // Cyan to green
          r = 0;
          g = 255;
          b = 255 - (value - 128) * 4;
          a = 255;
        } else {
          // Green to yellow to red
          r = (value - 192) * 4;
          g = 255;
          b = 0;
          a = 255;
        }

        const index = ((y * rect.width) + x) * 4;
        imageData.data[index] = r;
        imageData.data[index + 1] = g;
        imageData.data[index + 2] = b;
        imageData.data[index + 3] = a;
      }

      // Draw the spectrogram
      ctx.putImageData(imageData, 0, 0);

      // Draw frequency labels
      ctx.fillStyle = 'rgba(255, 255, 255, 0.8)';
      ctx.font = '10px monospace';
      ctx.textAlign = 'left';
      ctx.fillText('20kHz', 5, 12);
      ctx.fillText('1kHz', 5, rect.height / 2);
      ctx.fillText('20Hz', 5, rect.height - 5);
    };

    draw();

    return () => {
      if (animationFrameRef.current) {
        cancelAnimationFrame(animationFrameRef.current);
      }
    };
  }, [analyserNode]);

  return (
    <div className={cn('w-full h-full bg-card/50 border-2 border-accent/50 rounded-lg p-2', className)}>
      <div className="text-[10px] font-bold text-accent uppercase tracking-wider mb-2">
        Spectrogram
      </div>
      <div className="w-full h-[calc(100%-24px)] rounded bg-muted/30">
        <canvas
          ref={canvasRef}
          className="w-full h-full rounded"
        />
      </div>
    </div>
  );
}