wisp/src/lib/audio/wav.ts

// Minimal, dependency-free WAV (RIFF/WAVE) codec for PCM audio.
//
// Supports decoding:
//   - 16-bit integer PCM        (fmt audioFormat === 1)
//   - 32-bit IEEE float PCM     (fmt audioFormat === 3)
// Chunks may appear in any order; unknown chunks are skipped. Channels are
// de-interleaved into one Float32Array per channel, with samples in [-1, 1].
//
// Encoding produces a mono, 16-bit PCM WAVE file.
//
// All multi-byte header fields are little-endian (per the WAVE spec). We use
// DataView so endianness is explicit and never relies on the host platform.

/** A decoded WAVE file: sample rate plus one Float32Array per channel. */
export interface DecodedWav {
  sampleRate: number;
  /** One entry per channel; each holds that channel's samples in [-1, 1]. */
  channelData: Float32Array[];
}

const WAVE_FORMAT_PCM = 1;
const WAVE_FORMAT_IEEE_FLOAT = 3;

/** Read a 4-byte ASCII tag (e.g. "RIFF") at the given byte offset. */
function readTag(view: DataView, offset: number): string {
  return (
    String.fromCharCode(view.getUint8(offset)) +
    String.fromCharCode(view.getUint8(offset + 1)) +
    String.fromCharCode(view.getUint8(offset + 2)) +
    String.fromCharCode(view.getUint8(offset + 3))
  );
}

/**
 * Parse a RIFF/WAVE byte buffer.
 *
 * @throws if the buffer is too small, is not a RIFF container, is not a WAVE
 *   form, lacks fmt/data chunks, or uses an unsupported sample format.
 */
export function decodeWav(bytes: Uint8Array): DecodedWav {
  // A valid WAVE needs at least the 12-byte RIFF/WAVE descriptor.
  if (bytes.byteLength < 12) {
    throw new Error('Invalid WAV: file too small to contain a RIFF header');
  }

  const view = new DataView(
    bytes.buffer,
    bytes.byteOffset,
    bytes.byteLength,
  );

  if (readTag(view, 0) !== 'RIFF') {
    throw new Error("Invalid WAV: missing 'RIFF' magic");
  }
  if (readTag(view, 8) !== 'WAVE') {
    throw new Error("Invalid WAV: not a 'WAVE' file");
  }

  let audioFormat = -1;
  let numChannels = 0;
  let sampleRate = 0;
  let bitsPerSample = 0;
  let dataOffset = -1;
  let dataLength = 0;
  let haveFmt = false;
  let haveData = false;

  // Walk the chunk list starting right after the 12-byte descriptor.
  let offset = 12;
  while (offset + 8 <= bytes.byteLength) {
    const chunkId = readTag(view, offset);
    const chunkSize = view.getUint32(offset + 4, true);
    const chunkBody = offset + 8;

    if (chunkId === 'fmt ') {
      if (chunkBody + 16 > bytes.byteLength) {
        throw new Error('Invalid WAV: truncated fmt chunk');
      }
      audioFormat = view.getUint16(chunkBody, true);
      numChannels = view.getUint16(chunkBody + 2, true);
      sampleRate = view.getUint32(chunkBody + 4, true);
      // bytes 8..12: byteRate, bytes 12..14: blockAlign (derived; ignored)
      bitsPerSample = view.getUint16(chunkBody + 14, true);
      haveFmt = true;
    } else if (chunkId === 'data') {
      dataOffset = chunkBody;
      // Clamp to the actual buffer in case the declared size overruns it.
      dataLength = Math.min(chunkSize, bytes.byteLength - chunkBody);
      haveData = true;
    }
    // Unknown chunks fall through and are skipped.

    // Chunks are word-aligned: an odd size is followed by a pad byte.
    let advance = chunkSize;
    if (advance % 2 === 1) advance += 1;
    offset = chunkBody + advance;
  }

  if (!haveFmt) {
    throw new Error("Invalid WAV: missing 'fmt ' chunk");
  }
  if (!haveData) {
    throw new Error("Invalid WAV: missing 'data' chunk");
  }
  if (numChannels < 1) {
    throw new Error(`Invalid WAV: bad channel count ${numChannels}`);
  }

  const channelData: Float32Array[] = [];

  if (audioFormat === WAVE_FORMAT_PCM && bitsPerSample === 16) {
    const bytesPerSample = 2;
    const frameSize = bytesPerSample * numChannels;
    const numFrames = Math.floor(dataLength / frameSize);
    for (let c = 0; c < numChannels; c++) {
      channelData.push(new Float32Array(numFrames));
    }
    for (let frame = 0; frame < numFrames; frame++) {
      const base = dataOffset + frame * frameSize;
      for (let c = 0; c < numChannels; c++) {
        const int16 = view.getInt16(base + c * bytesPerSample, true);
        // Asymmetric int16 range: divide negatives by 32768, positives by 32767.
        const f = int16 < 0 ? int16 / 0x8000 : int16 / 0x7fff;
        channelData[c]![frame] = f;
      }
    }
  } else if (
    audioFormat === WAVE_FORMAT_IEEE_FLOAT &&
    bitsPerSample === 32
  ) {
    const bytesPerSample = 4;
    const frameSize = bytesPerSample * numChannels;
    const numFrames = Math.floor(dataLength / frameSize);
    for (let c = 0; c < numChannels; c++) {
      channelData.push(new Float32Array(numFrames));
    }
    for (let frame = 0; frame < numFrames; frame++) {
      const base = dataOffset + frame * frameSize;
      for (let c = 0; c < numChannels; c++) {
        channelData[c]![frame] = view.getFloat32(
          base + c * bytesPerSample,
          true,
        );
      }
    }
  } else {
    throw new Error(
      `Unsupported WAV format: audioFormat=${audioFormat}, ` +
        `bitsPerSample=${bitsPerSample} (supported: 16-bit PCM, 32-bit float)`,
    );
  }

  return { sampleRate, channelData };
}

/**
 * Encode mono samples (in [-1, 1]) into a 16-bit PCM WAVE file.
 * Out-of-range samples are clamped before quantization.
 */
export function encodeWav(
  samples: Float32Array,
  sampleRate: number,
): Uint8Array {
  const numChannels = 1;
  const bitsPerSample = 16;
  const bytesPerSample = bitsPerSample / 8;
  const blockAlign = numChannels * bytesPerSample;
  const byteRate = sampleRate * blockAlign;
  const dataSize = samples.length * bytesPerSample;
  const headerSize = 44;
  const totalSize = headerSize + dataSize;

  const buffer = new ArrayBuffer(totalSize);
  const view = new DataView(buffer);

  const writeTag = (offset: number, tag: string): void => {
    for (let i = 0; i < tag.length; i++) {
      view.setUint8(offset + i, tag.charCodeAt(i));
    }
  };

  // RIFF descriptor
  writeTag(0, 'RIFF');
  view.setUint32(4, totalSize - 8, true); // chunkSize = file size minus first 8
  writeTag(8, 'WAVE');

  // fmt chunk (16 bytes of PCM format data)
  writeTag(12, 'fmt ');
  view.setUint32(16, 16, true); // fmt chunk body size
  view.setUint16(20, WAVE_FORMAT_PCM, true);
  view.setUint16(22, numChannels, true);
  view.setUint32(24, sampleRate, true);
  view.setUint32(28, byteRate, true);
  view.setUint16(32, blockAlign, true);
  view.setUint16(34, bitsPerSample, true);

  // data chunk
  writeTag(36, 'data');
  view.setUint32(40, dataSize, true);

  let offset = headerSize;
  for (let i = 0; i < samples.length; i++) {
    let s = samples[i]!;
    // Clamp to [-1, 1] to avoid integer overflow/wraparound on quantize.
    if (s > 1) s = 1;
    else if (s < -1) s = -1;
    // Match the asymmetric int16 range used when decoding.
    const int16 = s < 0 ? Math.round(s * 0x8000) : Math.round(s * 0x7fff);
    view.setInt16(offset, int16, true);
    offset += bytesPerSample;
  }

  return new Uint8Array(buffer);
}