import { uniq } from 'lodash-es';
import {
  createContext,
  useCallback,
  useEffect,
  useMemo,
  useRef,
  useState,
} from 'react';
import { useLocalStorage } from 'usehooks-ts';
import { v4 as uuidv4 } from 'uuid';

import { useStores } from '@/app/(root)/AppProviders';
import useClip, { fetchClip } from '@/hooks/useClip';
import { useApiClient } from '@/lib/apiClient';
import { Clip } from '@/state/clipStore';

import downloadAsZip from './downloadAsZip';
import { songSectionColors } from './getSections';
import { EditProject, EditTrack } from './types';
import useStateHistory from './useStateHistory';

export const StemGroupId = {
  Vocals: 'Vocals',
  Backing_Vocals: 'Backing_Vocals',
  Drums: 'Drums',
  Bass: 'Bass',
  Guitar: 'Guitar',
  Keyboard: 'Keyboard',
  Percussion: 'Percussion',
  Strings: 'Strings',
  Synth: 'Synth',
  FX: 'FX',
  Brass: 'Brass',
  Woodwinds: 'Woodwinds',
};

export const StemTypeId = {
  'Lead Vocals': 0,
  // 'Lead Vocal': 0,
  'Electric Guitar': 1,
  'Backing Vocals': 2,
  // 'Drum Kit': 3,
  Drums: 3,
  Bass: 4,
  Piano: 5,
  'Acoustic Guitar': 6,
  Percussion: 7,
  Strings: 8,
  // 'String Section': 8,
  'Synth Pad': 9,
  Synthesizer: 10,
  Organ: 11,
  'Synth Bass': 12,
  'Lead Electric Guitar': 13,
  'Synth Keys': 14,
  'Rhythm Electric Guitar': 15,
  'Electronic Drum Kit': 16,
  'Noise effects': 17,
  'Electric Piano': 18,
  'Arr. Electric Guitar': 19,
  'Brass section': 20,
  'Distorted Electric Guitar': 21,
  'Upright Bass': 22,
  'Synth Strings': 23,
  'Synth Lead': 24,
  'Rhythm Acoustic Guitar': 25,
  'Intro Count (Click + Key)': 26,
  Flute: 27,
  Arpeggiator: 28,
  Tambourine: 29,
  Trumpet: 30,
  Harp: 31,
  'Synth Voice': 32,
  Accordion: 33,
  Fiddle: 34,
  'Synth Brass': 35,
  Violin: 36,
  'Pedal Steel Guitar': 37,
  'Sound effects': 38,
  Brass: 39,
  'Brass Instruments': 39,
  'Digital Piano': 40,
  'Tenor Saxophone': 41,
  Mandolin: 42,
  Trombone: 43,
  Clarinet: 44,
  'French Horn': 45,
  Banjo: 46,
  'Lead Acoustic Guitar': 47,
  'Rhythm Electric Guitar (Arpeggio)': 48,
  Glockenspiel: 49,
  Shaker: 50,
  'Electric Bass': 51,
  Timpani: 52,
  Vibes: 53,
  Harmonica: 54,
  'Hand Clap': 55,
  'Guitar Synth': 56,
  'Lap Steel Guitar': 57,
  'Arr. Acoustic Guitar': 58,
  Cello: 59,
  'Electronic Percussion': 60,
  Woodwinds: 61,
  Oboe: 62,
  'Rhythm Acoustic Guitar (Arpeggio)': 63,
  Saxophone: 64,
  'Drums and Percussion': 65,
  'Female Backing Vocals': 66,
  Marimba: 67,
  'Male Backing Vocals': 68,
  'Alto Saxophone': 69,
  Bells: 70,
  'Pitched Percussion': 71,
  'Synth Ambiant': 72,
  Dobro: 73,
  Celesta: 74,
  Congas: 75,
  'Synth Flute': 76,
  Harpsichord: 77,
  'Baritone Saxophone': 78,
  'Orchestra Hit': 79,
  'Female Lead Vocal': 80,
  'Male Lead Vocal': 81,
  Keyboard: 82,
  'Wind Chimes': 83,
  'Wind Instruments': 84,
  'Double Bass': 85,
  'Acoustic Drum Kit': 86,
  'Orchestral Percussion': 87,
  Xylophone: 88,
  Ukulele: 89,
  Orchestra: 90,
  FX: 91,
  // 'Sound Effects': 91,
  Cymbals: 92,
  Tuba: 93,
  Clap: 94,
  'Latin Percussion': 95,
  'Intro Count (Click)': 96,
  Bassoon: 97,
  Triangle: 98,
  'Snap (Fingers)': 99,
  Sitar: 100,
  'Tubular Bell': 101,
  Bongos: 102,
  Cowbell: 103,
  Sample: 104,
  'Metal bars': 105,
  'Steel Drums': 106,
  Guitar: 107,
  Viola: 108,
  'Soprano Saxophone': 109,
  'Claves/Woodblock': 110,
  Whistle: 111,
  'English Horn': 112,
  'Jingle Bells': 113,
  'Thumb Piano': 114,
  Piccolo: 115,
  'Intro Count (Key)': 116,
  Bagpipes: 117,
  Voice: 118,
  'Bass Drum': 119,
  'Music Box': 120,
  Rhodes: 121,
  'Appalachian Dulcimer': 122,
  'Human Beatbox': 123,
  Flugelhorn: 124,
  'Fretless Bass': 125,
  Theremin: 126,
  Vocoder: 127,
};

export const getBestStemId = (stemGroup: keyof typeof StemGroupId) => {
  return {
    [StemGroupId.Vocals]: StemTypeId['Lead Vocals'],
    [StemGroupId.Drums]: StemTypeId.Drums,
    [StemGroupId.Bass]: StemTypeId.Bass,
    [StemGroupId.Guitar]: StemTypeId.Guitar,
    [StemGroupId.Keyboard]: StemTypeId.Piano,
    [StemGroupId.Strings]: StemTypeId.Strings,
    [StemGroupId.Woodwinds]: StemTypeId.Woodwinds,
    [StemGroupId.Brass]: StemTypeId.Brass,
    [StemGroupId.Keyboard]: StemTypeId.Piano,
    [StemGroupId.Percussion]: StemTypeId.Percussion,
    [StemGroupId.Synth]: StemTypeId.Synthesizer,
  }[stemGroup];
};

const stemCategoryColors = {
  guitar: songSectionColors.bridge,
  bass: songSectionColors.verse,
  drums: songSectionColors.chorus,
  vocals: songSectionColors.intro,
  strings: songSectionColors.bridge,
  woodwinds: songSectionColors.bridge,
  brass: songSectionColors.bridge,
  keys: songSectionColors.bridge,
};

export const stemColors: Record<keyof typeof StemTypeId, string> = {
  // Vocals
  'Lead Vocals': stemCategoryColors.vocals,
  // 'Lead Vocal': stemCategoryColors.vocals,
  'Backing Vocals': stemCategoryColors.vocals,
  'Female Backing Vocals': stemCategoryColors.vocals,
  'Male Backing Vocals': stemCategoryColors.vocals,
  'Female Lead Vocal': stemCategoryColors.vocals,
  'Male Lead Vocal': stemCategoryColors.vocals,
  'Synth Voice': stemCategoryColors.vocals,
  Voice: stemCategoryColors.vocals,
  Vocoder: stemCategoryColors.vocals,
  'Human Beatbox': stemCategoryColors.vocals,

  // Guitars
  'Electric Guitar': stemCategoryColors.guitar,
  'Acoustic Guitar': stemCategoryColors.guitar,
  'Lead Electric Guitar': stemCategoryColors.guitar,
  'Rhythm Electric Guitar': stemCategoryColors.guitar,
  'Arr. Electric Guitar': stemCategoryColors.guitar,
  'Distorted Electric Guitar': stemCategoryColors.guitar,
  'Rhythm Acoustic Guitar': stemCategoryColors.guitar,
  'Lead Acoustic Guitar': stemCategoryColors.guitar,
  'Rhythm Electric Guitar (Arpeggio)': stemCategoryColors.guitar,
  'Guitar Synth': stemCategoryColors.guitar,
  'Lap Steel Guitar': stemCategoryColors.guitar,
  'Arr. Acoustic Guitar': stemCategoryColors.guitar,
  'Rhythm Acoustic Guitar (Arpeggio)': stemCategoryColors.guitar,
  'Pedal Steel Guitar': stemCategoryColors.guitar,
  Guitar: stemCategoryColors.guitar,
  Dobro: stemCategoryColors.guitar,

  // Bass
  Bass: stemCategoryColors.bass,
  'Synth Bass': stemCategoryColors.bass,
  'Upright Bass': stemCategoryColors.bass,
  'Electric Bass': stemCategoryColors.bass,
  'Double Bass': stemCategoryColors.bass,
  'Fretless Bass': stemCategoryColors.bass,

  // Drums & Percussion
  Drums: stemCategoryColors.drums,
  // 'Drum Kit': stemCategoryColors.drums,
  Percussion: stemCategoryColors.drums,
  'Electronic Drum Kit': stemCategoryColors.drums,
  Tambourine: stemCategoryColors.drums,
  Shaker: stemCategoryColors.drums,
  'Hand Clap': stemCategoryColors.drums,
  'Electronic Percussion': stemCategoryColors.drums,
  'Drums and Percussion': stemCategoryColors.drums,
  Congas: stemCategoryColors.drums,
  'Orchestral Percussion': stemCategoryColors.drums,
  Cymbals: stemCategoryColors.drums,
  Clap: stemCategoryColors.drums,
  'Latin Percussion': stemCategoryColors.drums,
  Triangle: stemCategoryColors.drums,
  'Snap (Fingers)': stemCategoryColors.drums,
  Bongos: stemCategoryColors.drums,
  Cowbell: stemCategoryColors.drums,
  'Claves/Woodblock': stemCategoryColors.drums,
  'Bass Drum': stemCategoryColors.drums,
  'Acoustic Drum Kit': stemCategoryColors.drums,

  // Strings
  Strings: stemCategoryColors.strings,
  // 'String Section': stemCategoryColors.strings,
  'Synth Strings': stemCategoryColors.strings,
  Violin: stemCategoryColors.strings,
  Cello: stemCategoryColors.strings,
  Viola: stemCategoryColors.strings,
  Harp: stemCategoryColors.strings,
  Fiddle: stemCategoryColors.strings,
  Sitar: stemCategoryColors.strings,
  'Appalachian Dulcimer': stemCategoryColors.strings,
  Ukulele: stemCategoryColors.strings,
  Banjo: stemCategoryColors.strings,
  Mandolin: stemCategoryColors.strings,

  // Woodwinds
  Flute: stemCategoryColors.woodwinds,
  'Synth Flute': stemCategoryColors.woodwinds,
  Clarinet: stemCategoryColors.woodwinds,
  Oboe: stemCategoryColors.woodwinds,
  Woodwinds: stemCategoryColors.woodwinds,
  Saxophone: stemCategoryColors.woodwinds,
  'Alto Saxophone': stemCategoryColors.woodwinds,
  'Tenor Saxophone': stemCategoryColors.woodwinds,
  'Baritone Saxophone': stemCategoryColors.woodwinds,
  'Soprano Saxophone': stemCategoryColors.woodwinds,
  Bassoon: stemCategoryColors.woodwinds,
  Piccolo: stemCategoryColors.woodwinds,
  'Wind Instruments': stemCategoryColors.woodwinds,
  'English Horn': stemCategoryColors.woodwinds,

  // Brass
  Brass: stemCategoryColors.brass,

  'Brass section': stemCategoryColors.brass,
  'Brass Instruments': stemCategoryColors.brass,
  'Synth Brass': stemCategoryColors.brass,
  Trumpet: stemCategoryColors.brass,
  Trombone: stemCategoryColors.brass,
  'French Horn': stemCategoryColors.brass,
  Tuba: stemCategoryColors.brass,
  Flugelhorn: stemCategoryColors.brass,

  // Keys & Synths
  Piano: stemCategoryColors.keys,
  'Synth Pad': stemCategoryColors.keys,
  Synthesizer: stemCategoryColors.keys,
  Organ: stemCategoryColors.keys,
  'Synth Keys': stemCategoryColors.keys,
  'Electric Piano': stemCategoryColors.keys,
  'Digital Piano': stemCategoryColors.keys,
  'Synth Lead': stemCategoryColors.keys,
  Arpeggiator: stemCategoryColors.keys,
  'Synth Ambiant': stemCategoryColors.keys,
  Keyboard: stemCategoryColors.keys,
  Harpsichord: stemCategoryColors.keys,
  Celesta: stemCategoryColors.keys,
  'Music Box': stemCategoryColors.keys,
  Rhodes: stemCategoryColors.keys,
  Theremin: stemCategoryColors.keys,
  Accordion: stemCategoryColors.keys,

  // Misc/Effects (using synth color as default)
  'Noise effects': stemCategoryColors.keys,
  'Sound effects': stemCategoryColors.keys,
  FX: stemCategoryColors.keys,
  // 'Sound Effects': stemCategoryColors.keys,
  'Orchestra Hit': stemCategoryColors.keys,
  Orchestra: stemCategoryColors.keys,
  Sample: stemCategoryColors.keys,
  'Intro Count (Click)': stemCategoryColors.keys,
  'Intro Count (Key)': stemCategoryColors.keys,
  'Intro Count (Click + Key)': stemCategoryColors.keys,

  // Pitched Percussion (using keys color)
  Glockenspiel: stemCategoryColors.keys,
  Vibes: stemCategoryColors.keys,
  Marimba: stemCategoryColors.keys,
  Bells: stemCategoryColors.keys,
  'Pitched Percussion': stemCategoryColors.keys,
  'Metal bars': stemCategoryColors.keys,
  'Steel Drums': stemCategoryColors.keys,
  'Tubular Bell': stemCategoryColors.keys,
  'Wind Chimes': stemCategoryColors.keys,
  Xylophone: stemCategoryColors.keys,
  'Jingle Bells': stemCategoryColors.keys,
  'Thumb Piano': stemCategoryColors.keys,
  Bagpipes: stemCategoryColors.keys,
  Timpani: stemCategoryColors.keys,
  Harmonica: stemCategoryColors.keys,
  Whistle: stemCategoryColors.keys,
};

export const getStemName = (clip: Clip) => {
  return (
    Object.entries(StemTypeId).find(
      (kv) => kv[1] === (clip?.metadata?.stem_type_id as any)
    )?.[0] || 'Unknown'
  );
};

export const smallTrackHeight = 40;
export const mediumTrackHeight = 80;
export const largeTrackHeight = 160;

export const useStemsContext = (clipId: string) => {
  const apiClient = useApiClient();

  const { clip } = useClip(clipId);
  const rootClipDuration = clip?.metadata?.duration || 0;

  const [memorizedState, setMemorizedState] = useLocalStorage(
    `memorized-state-${clipId}-march-20th-2025-v4`,
    null as EditProject | null
  );

  const {
    state: projectState,
    undo,
    redo,
    clearHistory: clearProjectStateHistory,
    setState: setProjectState,
  } = useStateHistory<EditProject>(
    () =>
      memorizedState || {
        tracks: [
          {
            id: uuidv4(),
            name: 'Full Song',
            amplitude: 1,
            balance: 0,
            solo: false,
            mute: false,
            stagedClips: [clipId],
            arrangedClips: [],
            currentVersionIndex: 0,
            height: mediumTrackHeight,
          },
        ],
      }
  );

  const fullSongTrack = projectState.tracks[0];
  const stemTracks = useMemo(
    () => projectState.tracks.slice(1),
    [projectState]
  );

  useEffect(() => {
    const onKeyDown = (e: KeyboardEvent) => {
      if (e.key === 'z' && (e.metaKey || e.ctrlKey)) {
        if (e.shiftKey) {
          redo();
        } else {
          undo();
        }
      } else if (e.key === 'y' && (e.metaKey || e.ctrlKey)) {
        redo();
      }
    };
    window.addEventListener('keydown', onKeyDown);
    return () => {
      window.removeEventListener('keydown', onKeyDown);
    };
  }, [projectState.tracks, clipId]);

  const setFullSongTrack = useCallback(
    (updateFn: (track: EditTrack) => EditTrack) => {
      setProjectState((prev) => ({
        ...prev,
        tracks: [
          updateFn(fullSongTrack),
          ...prev.tracks.filter((t) => t.id !== fullSongTrack.id),
        ],
      }));
    },
    [fullSongTrack]
  );

  const hasInitializedRef = useRef(false);
  useEffect(() => {
    if (!hasInitializedRef.current && memorizedState) {
      setProjectState(memorizedState);
      clearProjectStateHistory();
      hasInitializedRef.current = true;
    } else {
      setMemorizedState(projectState);
    }
    if (
      clip &&
      rootClipDuration &&
      !(
        fullSongTrack.arrangedClips.length &&
        fullSongTrack.stagedClips.includes(clipId)
      )
    ) {
      setFullSongTrack((prev) => ({
        ...prev,
        arrangedClips: [
          {
            id: uuidv4(),
            clipId: clipId,
            title: clip.title || '',
            contentSeconds: rootClipDuration,
            startSeconds: 0,
            endSeconds: rootClipDuration,
            readStartSeconds: 0,
          },
        ],
        stagedClips: [clipId],
      }));
      if (!hasInitializedRef.current) {
        hasInitializedRef.current = true;
        clearProjectStateHistory();
      }
    }
  }, [
    clip,
    clipId,
    rootClipDuration,
    fullSongTrack,
    clearProjectStateHistory,
    projectState,
  ]);

  const [selectedTrackId, setSelectedTrackId] = useState<string | null>(null);

  const updateTrack = useCallback(
    (trackId: string, updateFn: (track: EditTrack) => EditTrack) => {
      setProjectState((prev) => ({
        ...prev,
        tracks: prev.tracks.map((t) => (t.id === trackId ? updateFn(t) : t)),
      }));
    },
    [fullSongTrack]
  );

  const addStemTrack = useCallback(
    (name: string, stemType?: keyof typeof StemTypeId): EditTrack => {
      const newTrack: EditTrack = {
        id: uuidv4(),
        name,
        stemType,
        arrangedClips: [],
        stagedClips: [],
        height: mediumTrackHeight,
        amplitude: 1,
        balance: 0,
        solo: false,
        mute: false,
      };
      setProjectState((prev) => ({
        ...prev,
        tracks: [...prev.tracks, newTrack],
      }));
      return newTrack;
    },
    []
  );

  const deleteStemTrack = useCallback(
    (trackId: string) => {
      if (selectedTrackId === trackId) {
        setSelectedTrackId(null);
      }
      setProjectState((prev) => ({
        ...prev,
        tracks: prev.tracks.filter((t) => t.id !== trackId),
      }));
    },
    [selectedTrackId]
  );

  const duplicateStemTrack = useCallback((trackId: string) => {
    setProjectState((prev) => {
      const newTracks: EditTrack[] = [];
      for (const track of prev.tracks) {
        if (track.id === trackId) {
          newTracks.push({ ...track, id: uuidv4() });
        }
        newTracks.push(track);
      }
      return {
        ...prev,
        tracks: newTracks,
      };
    });
  }, []);

  const toggleTrackSolo = useCallback(
    (trackId: string, isShiftClick: boolean) => {
      if (isShiftClick) {
        updateTrack(trackId, (prev) => ({
          ...prev,
          solo: !prev.solo,
        }));
      } else {
        setProjectState((prev) => {
          const wasSoloed = prev.tracks.find((t) => t.id === trackId)?.solo;
          if (wasSoloed) {
            return {
              ...prev,
              tracks: prev.tracks.map((t) =>
                t.id === trackId ? { ...t, solo: false } : t
              ),
            };
          } else {
            return {
              ...prev,
              tracks: prev.tracks.map((t) =>
                t.id === trackId
                  ? { ...t, solo: true }
                  : t.solo
                    ? { ...t, solo: false }
                    : t
              ),
            };
          }
        });
      }
    },
    []
  );

  const { clips: clipStore } = useStores();

  const generateStems = useCallback(async () => {
    // deprecated function. returning empty array to keep this PR tight while still building
    return [] as Clip[];
  }, []);

  const generateStemsForTrack = useCallback(async () => {
    return [] as Clip[];
  }, []);

  const downloadZip = useCallback(async () => {
    const clips: Clip[] = [];

    console.log('downloading zip!');
    const allArrangedClipIDs = uniq(
      projectState.tracks
        .map((t) => t.arrangedClips.map((c) => c.clipId))
        .flat()
    );
    await Promise.all(
      allArrangedClipIDs.map(async (c) => {
        if (clipStore.clipById[c]) {
          clips.push(clipStore.clipById[c]);
        } else {
          clips.push(await fetchClip(clipStore, c, true));
        }
      })
    );

    await downloadAsZip(
      apiClient,
      clips.map((c) => clipStore.clipById[c.id]),
      clip?.title || 'Stems'
    );
  }, [clip, clipStore.clipById, projectState.tracks]);

  const anyTrackSolo = useMemo(() => {
    return stemTracks.some((t) => t.solo) || fullSongTrack.solo;
  }, [stemTracks, fullSongTrack]);

  const reorderTracks = useCallback(
    (trackId: string, targetTrackId: string, placeBefore: boolean) => {
      if (targetTrackId === fullSongTrack.id) return;
      if (trackId === fullSongTrack.id) return;
      setProjectState((prev) => {
        const trackIndex = prev.tracks.findIndex((t) => t.id === trackId);
        const targetTrackIndex = prev.tracks.findIndex(
          (t) => t.id === targetTrackId
        );
        const newTracks = [...prev.tracks];
        const trackToMove = newTracks[trackIndex];
        newTracks.splice(trackIndex, 1);
        if (placeBefore) {
          newTracks.splice(targetTrackIndex, 0, trackToMove);
        } else {
          newTracks.splice(targetTrackIndex + 1, 0, trackToMove);
        }
        return {
          ...prev,
          tracks: newTracks,
        };
      });
    },
    [fullSongTrack]
  );

  return useMemo(
    () => ({
      generateStems,
      fullSongTrack,
      setFullSongTrack,
      stemTracks,
      deleteStemTrack,
      duplicateStemTrack,
      updateTrack,
      addStemTrack,
      generateStemsForTrack,
      selectedTrackId,
      setSelectedTrackId,
      toggleTrackSolo,
      downloadZip,
      anyTrackSolo,
      reorderTracks,
    }),
    [
      generateStems,
      fullSongTrack,
      setFullSongTrack,
      stemTracks,
      deleteStemTrack,
      duplicateStemTrack,
      updateTrack,
      addStemTrack,
      generateStemsForTrack,
      selectedTrackId,
      setSelectedTrackId,
      toggleTrackSolo,
      downloadZip,
      anyTrackSolo,
      reorderTracks,
    ]
  );
};

const StemsContext = createContext<ReturnType<typeof useStemsContext>>(
  undefined as never
);

export default StemsContext;