const WAV_HEADER_SIZE = 44;
const IEEE_FLOAT_FORMAT = 3;
const BIT_DEPTH = 32;

export function encodeWAV(buffer: AudioBuffer): ArrayBuffer {
  const numChannels = buffer.numberOfChannels;
  const sampleRate = buffer.sampleRate;
  const bytesPerSample = BIT_DEPTH / 8;
  const blockAlign = numChannels * bytesPerSample;

  const channelData: Float32Array[] = [];
  for (let i = 0; i < numChannels; i++) {
    const data = new Float32Array(buffer.length);
    buffer.copyFromChannel(data, i);
    channelData.push(data);
  }

  const dataLength = buffer.length * numChannels * bytesPerSample;
  const bufferLength = WAV_HEADER_SIZE + dataLength;
  const arrayBuffer = new ArrayBuffer(bufferLength);
  const view = new DataView(arrayBuffer);

  writeWAVHeader(view, numChannels, sampleRate, blockAlign, dataLength);
  writePCMData(view, channelData, WAV_HEADER_SIZE);

  return arrayBuffer;
}

function writeWAVHeader(
  view: DataView,
  numChannels: number,
  sampleRate: number,
  blockAlign: number,
  dataLength: number
): void {
  const writeString = (offset: number, string: string) => {
    for (let i = 0; i < string.length; i++) {
      view.setUint8(offset + i, string.charCodeAt(i));
    }
  };

  writeString(0, 'RIFF');
  view.setUint32(4, 36 + dataLength, true);
  writeString(8, 'WAVE');
  writeString(12, 'fmt ');
  view.setUint32(16, 16, true);
  view.setUint16(20, IEEE_FLOAT_FORMAT, true);
  view.setUint16(22, numChannels, true);
  view.setUint32(24, sampleRate, true);
  view.setUint32(28, sampleRate * blockAlign, true);
  view.setUint16(32, blockAlign, true);
  view.setUint16(34, BIT_DEPTH, true);
  writeString(36, 'data');
  view.setUint32(40, dataLength, true);
}

function writePCMData(view: DataView, channelData: Float32Array[], startOffset: number): void {
  const numChannels = channelData.length;
  const numSamples = channelData[0].length;
  let offset = startOffset;

  for (let i = 0; i < numSamples; i++) {
    for (let channel = 0; channel < numChannels; channel++) {
      const sample = Math.max(-1, Math.min(1, channelData[channel][i]));
      view.setFloat32(offset, sample, true);
      offset += 4;
    }
  }
}

export function downloadWAV(buffer: AudioBuffer, filename: string = 'sound.wav'): void {
  const wav = encodeWAV(buffer);
  const blob = new Blob([wav], { type: 'audio/wav' });
  const url = URL.createObjectURL(blob);

  const link = document.createElement('a');
  link.href = url;
  link.download = filename;
  link.click();

  URL.revokeObjectURL(url);
}
