ep-133-project-export/src/lib/exporters/utils.ts

import * as Sentry from '@sentry/react';
import { unzippedBackupAtom } from '~/atoms/droppedProjectFile';
import { ExportStatus, ProjectRawData, SoundInfo } from '../../types/types';
import { getFile, getFileNodeByPath } from '../midi/fs';
import { pcmToWavBlob } from '../pcmToWav';
import { store } from '../store';
import { AbortError, audioFormatAsBitDepth, findSoundIdByPad } from '../utils';

let _colorIndex = 0;

export const TICKS_PER_BEAT = 96; // (384 / 4 beats)
const COLORS = [
  '#B2B0E8',
  '#E27D60',
  '#26A693',
  '#E8A87C',
  '#85C1A9',
  '#C38D9E',
  '#41B3A3',
  '#F2B880',
  '#7DAF9C',
  '#F47261',
  '#9D6A89',
  '#5AA9A4',
  '#7A2A80',
  '#8FB996',
  '#47B267',
  '#B089A3',
  '#6CA6A3',
  '#A0C4B0',
  '#F5A97F',
  '#BBA0C0',
  '#79B2B2',
];

function extractPcmFromWav(wavData: Uint8Array): Uint8Array {
  const view = new DataView(wavData.buffer, wavData.byteOffset, wavData.byteLength);

  const riff = String.fromCharCode(wavData[0], wavData[1], wavData[2], wavData[3]);
  if (riff !== 'RIFF') {
    throw new Error('Invalid WAV file: missing RIFF header');
  }

  const wave = String.fromCharCode(wavData[8], wavData[9], wavData[10], wavData[11]);
  if (wave !== 'WAVE') {
    throw new Error('Invalid WAV file: missing WAVE format');
  }

  let offset = 12;
  while (offset < wavData.length - 8) {
    const chunkId = String.fromCharCode(
      wavData[offset],
      wavData[offset + 1],
      wavData[offset + 2],
      wavData[offset + 3],
    );
    const chunkSize = view.getUint32(offset + 4, true);

    if (chunkId === 'data') {
      return wavData.slice(offset + 8, offset + 8 + chunkSize);
    }

    offset += 8 + chunkSize;
    if (chunkSize % 2 !== 0) {
      offset += 1;
    }
  }

  throw new Error('Invalid WAV file: data chunk not found');
}

export function parseWavMetadata(wavData: Uint8Array) {
  const view = new DataView(wavData.buffer, wavData.byteOffset, wavData.byteLength);

  const riff = String.fromCharCode(wavData[0], wavData[1], wavData[2], wavData[3]);
  if (riff !== 'RIFF') {
    throw new Error('Invalid WAV file: missing RIFF header');
  }

  const wave = String.fromCharCode(wavData[8], wavData[9], wavData[10], wavData[11]);
  if (wave !== 'WAVE') {
    throw new Error('Invalid WAV file: missing WAVE format');
  }

  let offset = 12;
  let channels = 1;
  let samplerate = 44100;
  let bitsPerSample = 16;
  let audioFormat = 1;
  let rootNote = 60;
  let teMeta: Record<string, unknown> | null = null;

  while (offset < wavData.length - 8) {
    const chunkId = String.fromCharCode(
      wavData[offset],
      wavData[offset + 1],
      wavData[offset + 2],
      wavData[offset + 3],
    );
    const chunkSize = view.getUint32(offset + 4, true);

    if (chunkId === 'fmt ') {
      audioFormat = view.getUint16(offset + 8, true);
      channels = view.getUint16(offset + 10, true);
      samplerate = view.getUint32(offset + 12, true);
      bitsPerSample = view.getUint16(offset + 22, true);
    }

    if (chunkId === 'smpl') {
      rootNote = view.getUint32(offset + 20, true);
    }

    if (chunkId === 'LIST') {
      const listData = wavData.subarray(offset + 8, offset + 8 + chunkSize);
      const listType = String.fromCharCode(listData[0], listData[1], listData[2], listData[3]);
      if (listType === 'INFO') {
        let subOffset = 4;
        while (subOffset < listData.length - 8) {
          const subId = String.fromCharCode(
            listData[subOffset],
            listData[subOffset + 1],
            listData[subOffset + 2],
            listData[subOffset + 3],
          );
          const subSize = new DataView(
            listData.buffer,
            listData.byteOffset + subOffset + 4,
            4,
          ).getUint32(0, true);

          if (subId === 'TNGE') {
            const raw = listData.subarray(subOffset + 8, subOffset + 8 + subSize);
            const jsonStr = new TextDecoder('ascii')
              .decode(raw)
              .replace(/\0/g, '')
              .trim();
            try {
              teMeta = JSON.parse(jsonStr);
            } catch {
              console.warn('Failed to parse TNGE metadata JSON:', jsonStr);
            }
            break;
          }

          subOffset += 8 + subSize;
          if (subSize % 2 !== 0) {
            subOffset += 1;
          }
        }
      }
    }

    offset += 8 + chunkSize;
    if (chunkSize % 2 !== 0) {
      offset += 1;
    }
  }

  let format: 's16' | 's24' | 'float';
  if (audioFormat === 3 && bitsPerSample === 32) {
    format = 'float';
  } else if (bitsPerSample === 24) {
    format = 's24';
  } else {
    format = 's16';
  }

  return {
    channels,
    samplerate,
    format,
    rootNote: (teMeta?.['sound.rootnote'] as number) ?? rootNote,
    teMeta,
  };
}

export async function downloadPcm(
  soundId: number,
  progressCallback?: (bytesRead: number, totalRemaining: number) => void,
) {
  const backupZip = store.get(unzippedBackupAtom);

  if (backupZip) {
    const wavFile = Object.values(backupZip.files).find((file) =>
      file.name.startsWith(`/sounds/${String(soundId).padStart(3, '0')} `),
    );

    if (!wavFile) {
      throw new Error(`Sound file for sound ID ${soundId} not found in backup`);
    }

    const wavData = await wavFile.async('uint8array');
    return extractPcmFromWav(wavData);
  }

  const fileNode = await getFileNodeByPath(`/sounds/${String(soundId).padStart(3, '0')}.pcm`);

  if (!fileNode) {
    throw new Error(`Sound file for sound ID ${soundId} not found`);
  }

  const fileData = await getFile(fileNode.nodeId, progressCallback);

  return fileData.data;
}

export function getSampleName(name: string | undefined, soundId: number, extension = true) {
  if (soundId >= 1000) {
    return '';
  }

  if (soundId === 0 && !name) {
    return '';
  }

  const id = soundId.toString().padStart(3, '0');
  const n = name ? `${id} ${name}` : `${id} sample`;

  return extension ? `${n}.wav` : n;
}

function getSoundsInfoFromProject(data: ProjectRawData, exportAllPadsWithSamples = false) {
  const snds: SoundInfo[] = [];
  const existingSounds = new Set<number>();
  const usedSoundIds = new Set<number>();

  for (const scene of data.scenes) {
    for (const pattern of scene.patterns) {
      if (pattern.notes.length > 0) {
        const sid = findSoundIdByPad(pattern.pad, data.pads);
        if (sid) {
          usedSoundIds.add(sid);
        }
      }
    }
  }

  for (const group in data.pads) {
    for (const pad of data.pads[group]) {
      if (pad.soundId <= 0) {
        continue;
      }

      if (!exportAllPadsWithSamples && !usedSoundIds.has(pad.soundId)) {
        continue;
      }

      if (existingSounds.has(pad.soundId)) {
        continue;
      }

      const soundMeta = data.sounds.find((s) => s.id === pad.soundId);
      if (!soundMeta) {
        continue;
      }

      snds.push({ soundId: pad.soundId, soundMeta: soundMeta.meta });
      existingSounds.add(pad.soundId);
    }
  }

  return snds;
}

export async function collectSamples(
  data: ProjectRawData,
  progressCallback: ({ progress, status }: ExportStatus) => void,
  abortSignal: AbortSignal,
  exportAllPadsWithSamples = false,
) {
  const projectSounds = getSoundsInfoFromProject(data, exportAllPadsWithSamples);

  const samples: { name: string; data: Blob }[] = [];
  const downloaded: string[] = [];
  const missing: { name: string; error: string }[] = [];
  const percentStart = 3;
  const totalSounds = projectSounds.length || 1;
  const percentPerSound = 80 / totalSounds;
  let cnt = 0;

  for (const snd of projectSounds) {
    if (abortSignal.aborted) {
      throw new AbortError();
    }

    const fileName = getSampleName(snd.soundMeta.name, snd.soundId);

    try {
      progressCallback({
        progress: percentStart + percentPerSound * cnt,
        status: `Downloading sound: ${fileName}`,
      });

      const result = await downloadPcm(snd.soundId, (bytesRead, totalRemaining) => {
        if (abortSignal.aborted) {
          throw new AbortError();
        }

        const currentSoundProgress = bytesRead / (totalRemaining / percentPerSound);
        progressCallback({
          progress: percentStart + percentPerSound * cnt + currentSoundProgress,
          status: `Downloading sound: ${fileName}`,
        });
      });

      const wavBlob = pcmToWavBlob(
        result,
        snd.soundMeta.samplerate,
        audioFormatAsBitDepth(snd.soundMeta.format),
        snd.soundMeta.channels,
      );

      samples.push({
        name: fileName,
        data: wavBlob,
      });

      downloaded.push(fileName);
    } catch (err) {
      // only report when not aborted
      if (!abortSignal.aborted) {
        console.error(err);
        Sentry.captureException(err);
      }

      missing.push({
        name: fileName,
        error: err instanceof Error ? err.message : 'Unknown error',
      });
    } finally {
      cnt++;
    }
  }

  const sampleReport = { downloaded, missing };

  progressCallback({
    progress: projectSounds.length === 0 ? percentStart + 80 : percentStart + percentPerSound * cnt,
    status: 'Sample collection completed',
    sampleReport,
  });

  return { samples, sampleReport };
}

export function getNextColor() {
  const color = COLORS[_colorIndex];
  _colorIndex = (_colorIndex + 1) % COLORS.length;
  return color;
}

export function ticksToMidiTicks(ticks: number, ppq = 480) {
  return Math.round((ticks / TICKS_PER_BEAT) * ppq);
}

export function getQuarterNotesPerBar(numerator: number, denominator: number) {
  return (4 / denominator) * numerator;
}