import { Dispatch, SetStateAction, useMemo } from 'react';

import stemIconMap from '@/app/(root)/stems/stemIconMap';
import { ButtonShape } from '@/components/button/Button';
import {
  ContextMenuGroup,
  ContextMenuItem,
  ContextMenuTrigger,
} from '@/components/contextMenu/ContextMenu';
import { StemGroupId } from '@/components/edit2025/StemsContext';
import {
  InstrumentBackupVocalsIcon,
  InstrumentLeadVocalsIcon,
  MusicNoteIcon,
  StemAddIcon,
  TriangleDownIcon,
} from '@/icons';

import { InstrumentSpec } from '../types';
import { InstrumentPickerButton } from './components';

export const getInstrumentName = (instrument: InstrumentSpec) => {
  if (instrument.type === 'custom') {
    return instrument.prompt || 'New Track';
  } else if (instrument.type === 'nonVocalPreset') {
    return instrument.prompt.replaceAll('_', ' ');
  } else if (instrument.type === 'vocalPreset') {
    return instrument.prompt.replaceAll('_', ' ');
  } else {
    return 'New Track';
  }
};

export const isNewTrackOrCurrentInstrumentName = (
  name: string,
  instrument: InstrumentSpec
) => {
  return name === 'New Track' || name === getInstrumentName(instrument);
};

export const getInstrumentTypeIcon = (instrument: InstrumentSpec) => {
  if (instrument.type === 'song') {
    return <MusicNoteIcon />;
  } else if (instrument.type === 'vocalPreset') {
    if (instrument.prompt === 'Vocals') {
      return <InstrumentLeadVocalsIcon />;
    } else {
      return <InstrumentBackupVocalsIcon />;
    }
  } else if (instrument.type === 'custom') {
    return <StemAddIcon />;
  } else if (instrument.type === 'nonVocalPreset') {
    return (
      stemIconMap[instrument.prompt as keyof typeof stemIconMap]?.() ?? (
        <StemAddIcon />
      )
    );
  } else {
    throw new Error(`Invalid instrument type: ${JSON.stringify(instrument)}`);
  }
};

export function InstrumentPickerMenuContents({
  setInstrument,
}: {
  setInstrument: (instrument: InstrumentSpec) => void;
}) {
  return (
    <>
      <ContextMenuGroup>
        {Object.keys(StemGroupId)
          .filter((group) => !['Vocals', 'Backing_Vocals'].includes(group))
          .map((stemGroup) => (
            <ContextMenuItem
              key={stemGroup}
              icon={
                stemIconMap[stemGroup as keyof typeof stemIconMap]?.() ?? (
                  <StemAddIcon />
                )
              }
              onClick={() => {
                setInstrument({
                  type: 'nonVocalPreset',
                  prompt: stemGroup,
                });
              }}
            >
              {stemGroup.replace('_', ' ')}
            </ContextMenuItem>
          ))}
      </ContextMenuGroup>
      <ContextMenuGroup>
        {Object.keys(StemGroupId)
          .filter((group) => ['Vocals', 'Backing_Vocals'].includes(group))
          .map((stemGroup) => (
            <ContextMenuItem
              key={stemGroup}
              icon={
                stemIconMap[stemGroup as keyof typeof stemIconMap]?.() ?? (
                  <StemAddIcon />
                )
              }
              onClick={() => {
                setInstrument({
                  type: 'vocalPreset',
                  prompt: stemGroup,
                });
              }}
            >
              {stemGroup.replace('_', ' ')}
            </ContextMenuItem>
          ))}
      </ContextMenuGroup>
      <ContextMenuGroup>
        <ContextMenuItem
          icon={<MusicNoteIcon />}
          onClick={() => {
            setInstrument({
              type: 'song',
            });
          }}
        >
          Song
        </ContextMenuItem>
        <ContextMenuItem
          icon={<StemAddIcon />}
          onClick={() => {
            setInstrument({
              type: 'custom',
              prompt: '',
            });
          }}
        >
          Custom
        </ContextMenuItem>
      </ContextMenuGroup>
    </>
  );
}

export default function ContextualBarInstrumentPicker({
  instrument,
  setInstrument,
}: {
  instrument: InstrumentSpec;
  setInstrument: Dispatch<SetStateAction<InstrumentSpec>>;
}) {
  const instrumentTypeIcon = useMemo(() => {
    return getInstrumentTypeIcon(instrument);
  }, [instrument]);

  const instrumentTypeLabel = useMemo(() => {
    if (instrument.type === 'song') {
      return 'Song';
    } else if (instrument.type === 'vocalPreset') {
      return instrument.prompt.replaceAll('_', ' ');
    } else if (instrument.type === 'custom') {
      return 'Custom';
    } else if (instrument.type === 'nonVocalPreset') {
      return instrument.prompt.replaceAll('_', ' ');
    } else {
      throw new Error(`Invalid instrument type: ${JSON.stringify(instrument)}`);
    }
  }, [instrument]);

  return (
    <ContextMenuTrigger
      placement='top-right'
      ButtonComponent={(props) => (
        <InstrumentPickerButton
          shape={ButtonShape.Pill}
          iconStart={instrumentTypeIcon}
          iconEnd={<TriangleDownIcon className='-mx-1 h-4 w-4' />}
          {...props}
        >
          <span className='button-label'>{instrumentTypeLabel}</span>
        </InstrumentPickerButton>
      )}
      ContentsComponent={() => (
        <InstrumentPickerMenuContents setInstrument={setInstrument} />
      )}
    />
  );
}