import { omit } from 'lodash-es';
import { useCallback, useMemo, useState } from 'react';

import { toast } from '@/components/toast/Toast';
import { createContext } from '@/hooks/useContextSelector';
import audioContext from '@/lib/audioContext';
import { sleep } from '@/utils/utils';

export type AudioInputSpec = {
  id: string;
  index: number;
  label: string;
  numberOfChannels: number;
  isDefault: boolean;
};

export type AudioInputPreview = {
  id: string;
  index: number;
  label: string;
  isDefault: boolean;
  getNumberOfChannels: () => Promise<number>;
};

type SettingOverridesType = Partial<{
  echoCancellation: boolean;
  autoGainControl: boolean;
  noiseSuppression: boolean;
}>;

const startOneStream = async (
  deviceId: string,
  settingOverrides?: SettingOverridesType,
  channel?: number
) => {
  const mediaStream = await new Promise<MediaStream>((resolve, reject) => {
    navigator.mediaDevices
      .getUserMedia({
        audio: {
          deviceId: { exact: deviceId },
          echoCancellation: false,
          autoGainControl: false,
          noiseSuppression: false,
          ...settingOverrides,
          latency: 0,
          sampleRate: audioContext.sampleRate,
          channelCount: 2,
        } as MediaTrackConstraints, // typescript is missing the latency property
        video: false,
      })
      .then((inputStream) => {
        resolve(inputStream);
      })
      .catch((e) => {
        reject(e);
      });
  });

  if (channel === undefined) {
    return {
      routedStream: mediaStream,
      rootStream: mediaStream,
    };
  }

  const source = audioContext.createMediaStreamSource(mediaStream);
  const splitter = audioContext.createChannelSplitter(
    mediaStream.getAudioTracks()[0]?.getSettings()?.channelCount ?? 2
  );
  source.connect(splitter);
  const merger = audioContext.createChannelMerger(1);
  splitter.connect(merger, channel, 0);
  const destination = audioContext.createMediaStreamDestination();
  merger.connect(destination);
  return { routedStream: destination.stream, rootStream: mediaStream };
};

const cachedChannelCounts: Partial<Record<string, Promise<number>>> = {};

const fullInputCache: Partial<Record<string, Promise<AudioInputSpec>>> = {};

export const getFullInput = async (inputPreview: AudioInputPreview) => {
  if (!fullInputCache[inputPreview.id]) {
    fullInputCache[inputPreview.id] = new Promise(async (resolve) => {
      resolve({
        ...omit(inputPreview, ['getNumberOfChannels']),
        numberOfChannels: await inputPreview.getNumberOfChannels(),
      } as AudioInputSpec);
    });
  }

  return fullInputCache[inputPreview.id]!;
};

const inputStreamsRef: {
  [key: string]: Promise<{
    routedStream: MediaStream;
    rootStream: MediaStream;
  }>;
} = {};

const inputReferenceCounts: { [key: string]: number } = {};

const releaseInputStream = (id: string) => {
  inputReferenceCounts[id] = (inputReferenceCounts[id] || 0) - 1;
  if (inputReferenceCounts[id] <= 0) {
    delete inputReferenceCounts[id];
    inputStreamsRef[id].then(({ rootStream, routedStream }) => {
      rootStream.getTracks().forEach((track) => {
        track.stop();
      });
      routedStream.getTracks().forEach((track) => {
        track.stop();
      });
    });
    delete inputStreamsRef[id];
  }
};

const requestInputStream = (
  id: string,
  channel: number | undefined,
  callback: (stream: MediaStream) => void
) => {
  inputReferenceCounts[id] = (inputReferenceCounts[id] || 0) + 1;
  if (!inputStreamsRef[id]) {
    inputStreamsRef[id] = startOneStream(id, undefined, channel);
  }
  inputStreamsRef[id].then(({ routedStream }) => callback(routedStream));
  return () => releaseInputStream(id);
};

const makeInputChannelCountGetter = (id: string) => async () => {
  if (!cachedChannelCounts[id]) {
    cachedChannelCounts[id] = new Promise(async (resolve, reject) => {
      try {
        const release = requestInputStream(id, undefined, (stream) => {
          try {
            const tracks = stream.getAudioTracks();
            const numberOfChannels =
              tracks[0].getSettings().channelCount || tracks.length;
            release();
            resolve(numberOfChannels);
          } catch (e) {
            console.error('error counting tracks');
            reject(e);
          }
        });
      } catch (e) {
        console.error('error requesting stream input');
        reject(e);
      }
    });
  }

  return (await cachedChannelCounts[id])!;
};

const useAudioInputStream = () => {
  const [isDiscoveringInputs, setIsDiscoveringInputs] = useState(false);
  const [inputs, setInputs] = useState<{
    [key: string]: AudioInputPreview;
  } | null>(null);

  const discoverInputs: () => Promise<{ [key: string]: AudioInputPreview }> =
    useCallback(async () => {
      let mediaStream: MediaStream;
      try {
        setIsDiscoveringInputs(true);
        mediaStream = await navigator.mediaDevices.getUserMedia({
          audio: {
            echoCancellation: false,
            autoGainControl: false,
            noiseSuppression: false,
            latency: 0,
            sampleRate: audioContext.sampleRate,
          } as MediaTrackConstraints, // typescript is missing the latency property
          video: false,
        });
      } catch (e) {
        console.error(e);
        setIsDiscoveringInputs(false);
        toast({
          title: 'Error accessing audio inputs.',
          description:
            'You may need to grant microphone permission to suno.com through your browser settings.',
          status: 'error',
        });
        throw e;
      }

      try {
        const devices = await navigator.mediaDevices.enumerateDevices();

        const spec: { [key: string]: AudioInputPreview } = {};

        const inputDevices = devices.filter(
          ({ kind }) => kind === 'audioinput'
        ) as InputDeviceInfo[];

        for (let i = 0; i < inputDevices.length; i++) {
          const device = inputDevices[i];
          const labelWithoutDefault = device.label.replace('Default - ', '');
          if (
            Object.values(spec).some(
              (s) =>
                (s.label === labelWithoutDefault || s.label === device.label) &&
                (s.isDefault || i === 0)
            )
          ) {
            console.log(spec);
            continue;
          }
          spec[device.deviceId] = {
            index: i,
            id: device.deviceId,
            label: labelWithoutDefault,
            isDefault: i === 0,
            getNumberOfChannels: makeInputChannelCountGetter(device.deviceId),
          };
        }

        setInputs(spec);
        await sleep(1);
        mediaStream.getTracks().forEach((track) => {
          track.stop();
        });
        return spec;
      } catch (e) {
        console.error(e);
        toast({
          title: 'Error accessing audio inputs.',
          description: 'Please try again.',
          status: 'error',
        });
        throw e;
      } finally {
        setIsDiscoveringInputs(false);
      }
    }, []);

  return useMemo(
    () => ({ requestInputStream, discoverInputs, inputs, isDiscoveringInputs }),
    [discoverInputs, inputs, isDiscoveringInputs]
  );
};

const AudioInputStreamContext = createContext<
  ReturnType<typeof useAudioInputStream>
>(undefined as never);

export const AudioInputStreamContextProvider = ({
  children,
}: {
  children: React.ReactNode;
}) => {
  return (
    <AudioInputStreamContext.Provider value={useAudioInputStream()}>
      {children}
    </AudioInputStreamContext.Provider>
  );
};

export default AudioInputStreamContext;