import localforage from 'localforage';

import audioContext from '@/lib/audioContext';
import downloadAndDecodeAudio from '@/lib/downloadAndDecodeAudio';
import { sleep } from '@/utils/utils';

import { PojoAudioBuffer } from './types';

const ENABLE_INDEXED_DB = false;

export type AudioFormat = 'mp3' | 'opus';

const buffers: Record<string, PojoAudioBuffer> = {};

export const resampleAudioBuffer = async (
  audioBuffer: AudioBuffer,
  sampleRate: number
) => {
  const offlineContext = new OfflineAudioContext(
    audioBuffer.numberOfChannels,
    audioBuffer.duration * sampleRate,
    sampleRate
  );
  const offlineSource = offlineContext.createBufferSource();
  offlineSource.buffer = audioBuffer;
  offlineSource.connect(offlineContext.destination);
  offlineSource.start();
  return await offlineContext.startRendering();
};

const resamplePojoBuffer = async (
  pojoAudioBuffer: PojoAudioBuffer,
  sampleRate: number
) => {
  const audioBuffer = audioBufferFromPojoAudioBuffer(pojoAudioBuffer);
  return pojoAudioBufferFromAudioBuffer(
    await resampleAudioBuffer(audioBuffer, sampleRate)
  );
};

export const pojoAudioBufferFromAudioBuffer = (audioBuffer: AudioBuffer) => {
  const buffer: PojoAudioBuffer = {
    sampleRate: audioBuffer.sampleRate,
    channelDataList: [],
  };
  for (let i = 0; i < audioBuffer.numberOfChannels; i++) {
    buffer.channelDataList.push(audioBuffer.getChannelData(i));
  }
  return buffer;
};

export const audioBufferFromPojoAudioBuffer = (buffer: PojoAudioBuffer) => {
  const audioBuffer = audioContext.createBuffer(
    buffer.channelDataList.length,
    buffer.channelDataList[0].length,
    buffer.sampleRate
  );
  for (let i = 0; i < buffer.channelDataList.length; i++) {
    audioBuffer.getChannelData(i).set(buffer.channelDataList[i]);
  }
  return audioBuffer;
};

const savePojoBuffer = (id: string, buffer: PojoAudioBuffer) => {
  buffers[id] = buffer;
  if (ENABLE_INDEXED_DB) localforage.setItem(`buffer:${id}`, buffer);
};

export const getCachedPojoBuffer = (id: string, format: AudioFormat) => {
  return buffers[`${id}.${format}`];
};

const getBufferFromIndexedDB = async (id: string) => {
  if (buffers[id]) {
    return buffers[id];
  }

  let buffer = ENABLE_INDEXED_DB
    ? await localforage.getItem<PojoAudioBuffer>(`buffer:${id}`)
    : null;

  if (buffer) {
    if (buffer.sampleRate !== audioContext.sampleRate) {
      buffer = await resamplePojoBuffer(buffer, audioContext.sampleRate);
    }
    buffers[id] = buffer;
  }

  return buffer;
};

const requestsInFlight: Partial<
  Record<string, Promise<PojoAudioBuffer | null>>
> = {};

export const getPojoBuffer = async (
  id: string,
  format: AudioFormat = 'mp3'
) => {
  const cached = getCachedPojoBuffer(id, format);
  if (cached) {
    return cached;
  }

  const filename = `${id}.${format}`;

  if (!requestsInFlight[filename]) {
    requestsInFlight[filename] = (async () => {
      const maxRetries = 5;
      for (let retry = 0; retry < maxRetries; retry++) {
        try {
          const pojoBufferFromIndexedDB =
            await getBufferFromIndexedDB(filename);
          if (pojoBufferFromIndexedDB) {
            return pojoBufferFromIndexedDB;
          }
          const audioBuffer = await downloadAndDecodeAudio(
            `https://cdn1.suno.ai/${filename}`
          );
          const pojoBuffer = pojoAudioBufferFromAudioBuffer(audioBuffer);
          savePojoBuffer(filename, pojoBuffer);
          return pojoBuffer;
        } catch (e) {
          console.error(
            'Error fetching buffer, ' +
              (retry < maxRetries - 1
                ? 'waiting 1s before trying again'
                : 'giving up'),
            e
          );
          await sleep(1000);
        }
      }
      return null;
    })();
  }

  return await requestsInFlight[filename];
};