import {
  addBeatsToSeconds,
  getBeatsFromZero,
  getSecondsBetween,
} from '@suno/studiokit/timeMapping';
import { diffArrays } from 'diff';
import { isEqual } from 'lodash-es';
import React, {
  createContext,
  useCallback,
  useContext,
  useEffect,
  useMemo,
  useRef,
  useState,
} from 'react';
import { useInterval } from 'usehooks-ts';
import { v4 as uuidv4 } from 'uuid';

import { useStores } from '@/app/(root)/AppProviders';
import { fetchClip } from '@/hooks/useClip';
import useDebounceCallback from '@/hooks/useDebounceCallback';
import useDismount from '@/hooks/useDismount';
import audioContext from '@/lib/audioContext';
import { ContextType } from '@/logging/contextTypes';
import { eventLogger } from '@/utils/event-logger';
import { EventNames } from '@/utils/event-names';

import EditModeContext from './EditModeContext';
import PreviewClipContext from './PreviewClipContext';
import SelectionContext, { getEditCategory } from './SelectionContext';
import StemsContext from './StemsContext';
import {
  audioBufferFromPojoAudioBuffer,
  getCachedPojoBuffer,
  getPojoBuffer,
} from './bufferStorage';
import getAutomationValue from './getAutomationValue';
import EditTimingContext from './queryHooks/EditTimingContext';
import { ArrangedClip, AutomationPoint, EditTiming, EditTrack } from './types';

export type EditorPlaybackTrack = {
  id: string;
  amplitude: number;
  balance: number;
  clips: ArrangedClip[];
};

export type EditorPlaybackState = {
  tracks: EditorPlaybackTrack[];
  skipTimeStart: number;
  skipTimeEnd: number;
  amplitudeAutomation: AutomationPoint[];
  songEndSeconds: number;
  playbackEndSeconds: number;
  timing: EditTiming;
};

const songSessionIdsByClipId = {} as Record<
  string,
  {
    value: string;
    isNew: boolean;
  }
>;
const consumeSongSessionId = (arrangementId: string) => {
  if (!songSessionIdsByClipId[arrangementId]) {
    songSessionIdsByClipId[arrangementId] = {
      value: uuidv4(),
      isNew: true,
    };
  }
  return songSessionIdsByClipId[arrangementId].value;
};

const consumeIsNewSongSessionId = (arrangementId: string) => {
  const songSessionIdInfo = songSessionIdsByClipId[arrangementId];
  const wasNew = songSessionIdInfo.isNew;
  songSessionIdInfo.isNew = false;
  return wasNew;
};

const startTimesByArrangementId = {} as Record<string, number>;

export const useEditorStatePlayer = (state: EditorPlaybackState) => {
  const tearDownRef = useRef<() => void>(() => {});
  const [playing, setPlaying] = useState(false);
  const playingRef = useRef(false);
  const lastSeekTimeRef = useRef<number>(0);
  const lastReportedTimeRef = useRef(0);
  const lastReportedAtRef = useRef(Date.now());
  const lastStartedFromRef = useRef(0);
  const masterGainNodeRef = useRef<GainNode | null>(null);
  const [volume, setVolumeState] = useState(1);
  const { clips, session } = useStores();

  const gainNodesRef = useRef<Record<string, GainNode>>({});
  const panNodesRef = useRef<Record<string, StereoPannerNode>>({});

  const scheduledBuffersRef = useRef<
    Record<
      string,
      {
        startSeconds: number;
        endSeconds: number;
        readStartSeconds: number;
        clipId: string;
        sourceNode: AudioBufferSourceNode;
      }[]
    >
  >({});

  useEffect(() => {
    playingRef.current = playing;
  }, [playing]);

  const setVolume = useCallback(
    (newVolume: number, andCommit: boolean = true) => {
      if (andCommit) {
        setVolumeState(newVolume);
      }
      masterGainNodeRef.current?.gain.setValueAtTime(newVolume, 0);
    },
    []
  );

  // providing a continuously-updating `currentTime` value will either be low-resolution, or will nuke react.
  // instead, let components call this function as often as they need to to update UI outside of the react lifecycle.
  const getCurrentTime = useCallback(() => {
    let unskippedCurrentTime = lastStartedFromRef.current;
    if (playingRef.current) {
      const now = Date.now();

      // lastReportedAtRef.current = real-world-time that we last checked the internal playback time.
      // lastReportedTimeRef.current = the playback time that was reported when we checked.
      // the current time is lastReportedTime plus the amount of time since it was reported.
      unskippedCurrentTime =
        lastReportedTimeRef.current + (now - lastReportedAtRef.current) / 1000;
    }
    if (state.skipTimeEnd > state.skipTimeStart) {
      if (unskippedCurrentTime > state.skipTimeStart) {
        return unskippedCurrentTime + (state.skipTimeEnd - state.skipTimeStart);
      }
    }
    return unskippedCurrentTime;
  }, [state.skipTimeStart, state.skipTimeEnd]);

  const getTrackContainingClip = useCallback(
    (id: string) => {
      return state.tracks.find((t) => t.clips.find((c) => c.id === id));
    },
    [state.tracks]
  );

  const getClipById = useCallback(
    (id: string) => {
      return state.tracks.flatMap((t) => t.clips).find((c) => c.id === id);
    },
    [state.tracks]
  );

  const playCountsToIncrementRef = useRef<Record<string, boolean>>({});

  const debouncedIncrementPlayCounts = useDebounceCallback(() => {
    const clipIds = Object.keys(playCountsToIncrementRef.current);
    const averagePlayCount =
      clipIds.reduce(
        (acc, c) => acc + (clips.clipById[c]?.play_count || 0),
        0
      ) / clipIds.length;

    clips.apiClient.POST('/api/gen/bulk_increment_play_counts/v2', {
      body: {
        gen_ids: clipIds,
        sample_factor: Math.floor(1 + averagePlayCount / 100),
      },
    });
    playCountsToIncrementRef.current = {};
  }, 1200);

  const logPlaybackEvent = useCallback(
    (
      inputActionName:
        | 'PlaySong'
        | 'PauseSong'
        | 'SongEnd'
        | 'SeekProgressBarPauseSong'
        | 'SeekProgressBarPlaySong',
      affectedClips: ArrangedClip[],
      triggeringInteractionTime = getCurrentTime()
    ) => {
      const currentTime = getCurrentTime();

      if (inputActionName === 'PlaySong' && affectedClips.length > 0) {
        affectedClips.forEach((arrangedClip) => {
          playCountsToIncrementRef.current[arrangedClip.clipId] = true;
        });
        debouncedIncrementPlayCounts();
      }

      affectedClips.forEach(async (arrangedClip) => {
        let startTime: number | undefined;
        let endTime: number | undefined;
        if (['PlaySong', 'SeekProgressBarPlaySong'].includes(inputActionName)) {
          startTimesByArrangementId[arrangedClip.id] =
            currentTime +
            (arrangedClip.startSeconds - arrangedClip.readStartSeconds);
        } else if (
          ['SongEnd', 'PauseSong', 'SeekProgressBarPauseSong'].includes(
            inputActionName
          )
        ) {
          startTime = startTimesByArrangementId[arrangedClip.id];
          endTime = triggeringInteractionTime;
          delete startTimesByArrangementId[arrangedClip.id];
        }

        const clip = await fetchClip(clips, arrangedClip.clipId, false);
        const clipUserId = clip?.user_id;
        const track = getTrackContainingClip(arrangedClip.id);

        if (
          endTime !== undefined &&
          startTime !== undefined &&
          endTime < startTime
        ) {
          return;
        }

        const currentSongSessionId = consumeSongSessionId(arrangedClip.id);
        const actionName = consumeIsNewSongSessionId(arrangedClip.id)
          ? 'PlayNewSong'
          : inputActionName;

        const event = {
          songSessionId: currentSongSessionId,
          hasClip: true,
          songId: arrangedClip.clipId,
          contextId: undefined,
          contextType: ContextType.Studio,
          startTime,
          endTime,
          isPlaying: playingRef.current,
          playDuration:
            endTime !== undefined && startTime !== undefined
              ? endTime - startTime
              : undefined,
          isAudioElementNull: false,
          audioElementCurrentTime: triggeringInteractionTime,
          actionName,
          isUserSongOwner:
            session?.userId !== undefined && session?.userId === clipUserId,
          volume: track ? track.amplitude * 100 : undefined,
          clickSourceUrl: location.pathname,
          isAutoplayOn: false,
          isRepeatOn: false,
          userId: session?.userId,
          previousSongSessionId: null,
          actionIndex: -1,
          songLength: arrangedClip.endSeconds - arrangedClip.startSeconds,
        };

        eventLogger.segmentTrack(EventNames.audioPlayerEvent, event, session);
      });
    },
    [playing, getTrackContainingClip, getClipById]
  );

  const lastPlayingClipsRef = useRef<ArrangedClip[]>([]);

  const getCurrentlyPlayingClips = useCallback(() => {
    const currentTime = getCurrentTime();
    const audibleTracks = state.tracks.filter((t) => t.amplitude > 0);
    const playingClips = !playingRef.current
      ? []
      : audibleTracks.flatMap((t) =>
          t.clips.filter(
            (c) => c.startSeconds <= currentTime && c.endSeconds > currentTime
          )
        );
    return playingClips.sort((a, b) => a.id.localeCompare(b.id));
  }, [state.tracks]);

  const logPassivePlaybackChanges = useCallback(() => {
    const lastPlayingClips = lastPlayingClipsRef.current;
    const currentlyPlayingClips = getCurrentlyPlayingClips();

    const lastPlayingClipsById = lastPlayingClips.reduce(
      (acc, clip) => {
        acc[clip.id] = clip;
        return acc;
      },
      {} as Record<string, ArrangedClip>
    );

    const currentlyPlayingClipsById = currentlyPlayingClips.reduce(
      (acc, clip) => {
        acc[clip.id] = clip;
        return acc;
      },
      {} as Record<string, ArrangedClip>
    );

    const diff = diffArrays(
      lastPlayingClips.map((c) => c.id),
      currentlyPlayingClips.map((c) => c.id)
    );

    const removedClips = diff
      .filter((change) => change.removed)
      .map(({ value }) => value.map((id) => lastPlayingClipsById[id]))
      .flat()
      .filter(Boolean) as ArrangedClip[];

    const addedClips = diff
      .filter((change) => change.added)
      .map(({ value }) => value.map((id) => currentlyPlayingClipsById[id]))
      .flat()
      .filter(Boolean) as ArrangedClip[];

    logPlaybackEvent('PlaySong', addedClips);
    logPlaybackEvent(
      playingRef.current ? 'SongEnd' : 'PauseSong',
      removedClips
    );

    lastPlayingClipsRef.current = currentlyPlayingClips;
  }, [getCurrentlyPlayingClips, getClipById, logPlaybackEvent]);

  const logSeekPlaybackChanges = useCallback(
    (seekFromTime: number) => {
      const currentlyPlayingClips = getCurrentlyPlayingClips();
      logPlaybackEvent(
        'SeekProgressBarPauseSong',
        lastPlayingClipsRef.current,
        seekFromTime
      );
      logPlaybackEvent(
        'SeekProgressBarPlaySong',
        currentlyPlayingClips,
        seekFromTime
      );

      lastPlayingClipsRef.current = currentlyPlayingClips;
    },
    [getCurrentTime, getCurrentlyPlayingClips, logPlaybackEvent]
  );

  const logSongEndPlaybackChanges = useCallback(
    (seekFromTime: number) => {
      const currentlyPlayingClips = getCurrentlyPlayingClips();
      logPlaybackEvent('SongEnd', lastPlayingClipsRef.current, seekFromTime);
      lastPlayingClipsRef.current = currentlyPlayingClips;
    },
    [getCurrentTime, getCurrentlyPlayingClips, logPlaybackEvent]
  );

  useInterval(logPassivePlaybackChanges, 200);

  const audioBuffersByPojoBufferIdRef = useRef<Record<string, AudioBuffer>>({});
  const currentEndTimeRef = useRef(Infinity);
  const playCountRef = useRef(0);

  const setTrackAmplitude = useCallback((trackId: string, value: number) => {
    const gainNode = gainNodesRef.current[trackId];
    if (gainNode) {
      gainNode.gain.linearRampToValueAtTime(
        value,
        audioContext.currentTime + 0.01
      );
    }
  }, []);

  const setTrackBalance = useCallback((trackId: string, value: number) => {
    const panNode = panNodesRef.current[trackId];
    if (panNode) {
      panNode.pan.linearRampToValueAtTime(
        value,
        audioContext.currentTime + 0.01
      );
    }
  }, []);

  const setTrackClips = useCallback(
    (trackId: string, clips: ArrangedClip[]) => {
      if (!scheduledBuffersRef.current[trackId]) {
        scheduledBuffersRef.current[trackId] = [];
      }
      const unmatchedClips: ArrangedClip[] = [];
      const sourceNodesToStop: AudioBufferSourceNode[] = [];
      const bufferIndicesToRemove: number[] = [];

      for (let i = 0; i < scheduledBuffersRef.current[trackId].length; i++) {
        const entry = scheduledBuffersRef.current[trackId][i];
        if (
          !clips.find(
            (clip) =>
              clip.clipId === entry.clipId &&
              clip.startSeconds === entry.startSeconds &&
              clip.readStartSeconds === entry.readStartSeconds &&
              clip.endSeconds === entry.endSeconds
          )
        ) {
          bufferIndicesToRemove.push(i);
          sourceNodesToStop.push(entry.sourceNode);
        }
      }

      clips.forEach((clip) => {
        const entry = scheduledBuffersRef.current[trackId].find(
          (entry) =>
            entry.clipId === clip.clipId &&
            entry.startSeconds === clip.startSeconds &&
            entry.readStartSeconds === clip.readStartSeconds &&
            entry.endSeconds === clip.endSeconds
        );
        if (!entry) {
          unmatchedClips.push(clip);
        }
      });

      bufferIndicesToRemove.reverse().forEach((index) => {
        scheduledBuffersRef.current[trackId].splice(index, 1);
      });

      sourceNodesToStop.forEach((sourceNode) => {
        sourceNode.stop();
        sourceNode.disconnect();
      });

      unmatchedClips.forEach(async (clip) => {
        const endTime = clip.endSeconds;
        if (endTime < getCurrentTime()) {
          console.log(
            'returning early because first endTime < getCurrentTime()',
            endTime,
            getCurrentTime()
          );
          return;
        }

        const buffer = audioContext.createBufferSource();
        const pojoBuffer = await getPojoBuffer(clip.clipId, 'mp3');
        if (!pojoBuffer || !playingRef.current) {
          console.log(
            'returning early because !pojoBuffer || !playingRef.current',
            endTime,
            getCurrentTime()
          );
          return;
        }
        buffer.buffer = audioBufferFromPojoAudioBuffer(pojoBuffer);

        if (endTime < getCurrentTime()) {
          console.log(
            'returning early because second endTime < getCurrentTime()',
            endTime,
            getCurrentTime()
          );
          return;
        }

        const gainNode = gainNodesRef.current[trackId];
        if (!gainNode) {
          console.log(
            'returning early because !gainNode',
            endTime,
            getCurrentTime()
          );
          return;
        }

        buffer.connect(gainNode);
        const currentTime = getCurrentTime();

        console.log(
          'gonna play!',
          clip.startSeconds,
          clip.endSeconds,
          clip.readStartSeconds,
          currentTime
        );

        buffer.start(
          audioContext.currentTime +
            Math.max(0, clip.startSeconds - currentTime),
          Math.max(0, currentTime - clip.startSeconds) + clip.readStartSeconds,
          clip.endSeconds - Math.max(currentTime, clip.startSeconds)
        );
        scheduledBuffersRef.current[trackId].push({
          startSeconds: clip.startSeconds,
          endSeconds: clip.endSeconds,
          readStartSeconds: clip.readStartSeconds,
          clipId: clip.clipId,
          sourceNode: buffer,
        });
      });
    },
    [getCurrentTime]
  );

  const addTrack = useCallback(
    (trackId: string, audioCtx: BaseAudioContext = audioContext) => {
      if (!masterGainNodeRef.current) return;
      gainNodesRef.current[trackId] = audioCtx.createGain();
      panNodesRef.current[trackId] = audioCtx.createStereoPanner();
      gainNodesRef.current[trackId].connect(panNodesRef.current[trackId]);
      panNodesRef.current[trackId].connect(masterGainNodeRef.current);
    },
    [setTrackClips]
  );

  const deleteTrack = useCallback((trackId: string) => {
    const gainNode = gainNodesRef.current[trackId];
    if (gainNode) {
      gainNode.disconnect();
      delete gainNodesRef.current[trackId];
    }
    const panNode = panNodesRef.current[trackId];
    if (panNode) {
      panNode.disconnect();
      delete panNodesRef.current[trackId];
    }
    if (scheduledBuffersRef.current[trackId]) {
      scheduledBuffersRef.current[trackId].forEach((entry) => {
        entry.sourceNode.stop();
        entry.sourceNode.disconnect();
      });
      delete scheduledBuffersRef.current[trackId];
    }
  }, []);

  const lastSeekFromTime = useRef(null as number | null);

  const play = useCallback(
    async (
      startTime: number = getCurrentTime(),
      endTime: number = state.playbackEndSeconds,
      audioCtx: BaseAudioContext = audioContext
    ) => {
      const isRealtimePlayback = audioCtx === audioContext;
      if (startTime > state.skipTimeEnd) {
        startTime -= state.skipTimeEnd - state.skipTimeStart;
      }
      const effectiveEndTime =
        endTime <= startTime ? state.songEndSeconds : endTime;
      currentEndTimeRef.current = effectiveEndTime;

      const startTimeBeats = getBeatsFromZero(startTime, state.timing);
      playCountRef.current++;
      const playCount = playCountRef.current;

      tearDownRef.current();
      const masterGainNode = audioCtx.createGain();
      const fadeGainNode = audioCtx.createGain();

      await Promise.all(
        state.tracks
          .map((track) =>
            track.clips.map((clip) => getPojoBuffer(clip.clipId, 'mp3'))
          )
          .flat()
      );

      if (playCount !== playCountRef.current) return;

      state.tracks.forEach((track) => {
        const gainNode = audioCtx.createGain();

        const effectiveClips: ArrangedClip[] = [];

        track.clips.forEach((originalClip) => {
          let clip = { ...originalClip };
          const fullyOutsideSkipRegion =
            (clip.startSeconds <= state.skipTimeStart &&
              clip.endSeconds <= state.skipTimeStart) ||
            (clip.startSeconds >= state.skipTimeEnd &&
              clip.endSeconds >= state.skipTimeEnd);

          if (fullyOutsideSkipRegion) {
            effectiveClips.push(clip);
            return;
          }
          if (
            clip.startSeconds < state.skipTimeStart &&
            clip.endSeconds > state.skipTimeStart
          ) {
            clip.endSeconds = state.skipTimeStart;
            effectiveClips.push(clip);
            clip = { ...originalClip };
          }

          if (
            clip.startSeconds < state.skipTimeEnd &&
            clip.endSeconds > state.skipTimeEnd
          ) {
            clip.readStartSeconds += state.skipTimeEnd - clip.startSeconds;
            clip.startSeconds = state.skipTimeEnd;
            effectiveClips.push(clip);
            clip = { ...originalClip };
          }

          if (fullyOutsideSkipRegion) {
            effectiveClips.push(clip);
          }
        });

        effectiveClips.forEach((clip) => {
          if (clip.startSeconds >= state.skipTimeEnd) {
            clip.startSeconds -= state.skipTimeEnd - state.skipTimeStart;
            clip.endSeconds -= state.skipTimeEnd - state.skipTimeStart;
          }
        });

        effectiveClips.forEach((clip) => {
          if (clip.endSeconds < startTime) return;
          const playbackDuration =
            Math.min(clip.endSeconds, effectiveEndTime) -
            Math.max(startTime, clip.startSeconds);
          if (playbackDuration <= 0) return;
          const cachedBuffer = getCachedPojoBuffer(clip.clipId, 'mp3');
          if (!cachedBuffer) return;
          const source = audioCtx.createBufferSource();
          source.buffer =
            audioBuffersByPojoBufferIdRef.current[clip.clipId] ||
            audioBufferFromPojoAudioBuffer(
              getCachedPojoBuffer(clip.clipId, 'mp3')
            );
          audioBuffersByPojoBufferIdRef.current[clip.clipId] = source.buffer;

          if (!scheduledBuffersRef.current[track.id]) {
            scheduledBuffersRef.current[track.id] = [];
          }

          scheduledBuffersRef.current[track.id].push({
            startSeconds: clip.startSeconds,
            endSeconds: clip.endSeconds,
            readStartSeconds: clip.readStartSeconds,
            clipId: clip.clipId,
            sourceNode: source,
          });

          source.connect(gainNode);
          const timeUntilClipStart = Math.max(0, clip.startSeconds - startTime);
          const clipContentStart = Math.max(
            clip.readStartSeconds,
            clip.readStartSeconds + (startTime - clip.startSeconds)
          );
          source.start(
            audioCtx.currentTime + timeUntilClipStart,
            clipContentStart,
            playbackDuration
          );
        });

        if (audioCtx instanceof AudioContext) {
          setPlaying(true);
        }

        const panNode = audioCtx.createStereoPanner();

        gainNode.connect(panNode);
        gainNode.gain.value = track.amplitude;
        gainNodesRef.current[track.id] = gainNode;

        panNode.connect(masterGainNode);
        panNode.pan.value = track.balance;
        panNodesRef.current[track.id] = panNode;
      });

      fadeGainNode.gain.value = getAutomationValue(
        state.amplitudeAutomation,
        startTimeBeats,
        1
      );

      state.amplitudeAutomation.forEach((point) => {
        if (point.beats <= startTimeBeats) return;
        const timeUntilPoint = getSecondsBetween(
          startTimeBeats,
          point.beats,
          state.timing
        );
        fadeGainNode.gain.linearRampToValueAtTime(
          point.value,
          audioCtx.currentTime + timeUntilPoint
        );
      });

      masterGainNodeRef.current = masterGainNode;

      masterGainNode.gain.value = volume;
      masterGainNode.connect(fadeGainNode);
      fadeGainNode.connect(audioCtx.destination);

      lastStartedFromRef.current = startTime;
      lastReportedTimeRef.current = startTime;
      lastReportedAtRef.current = Date.now();

      if (audioCtx instanceof AudioContext && audioCtx.state !== 'running') {
        audioCtx.resume();
      }

      const timeToWait = effectiveEndTime - startTime;

      // :aaaaaa:
      let stopTimeout: NodeJS.Timeout | undefined;
      if (Number.isFinite(timeToWait) && isRealtimePlayback) {
        stopTimeout = setTimeout(() => {
          logSongEndPlaybackChanges(getCurrentTime());
          setPlaying(false);
        }, timeToWait * 1000);
      }

      tearDownRef.current = () => {
        if (stopTimeout) {
          clearTimeout(stopTimeout);
        }

        fadeGainNode.disconnect();

        if (masterGainNodeRef.current) {
          masterGainNodeRef.current.disconnect();
          masterGainNodeRef.current = null;
        }

        if (gainNodesRef.current) {
          Object.values(gainNodesRef.current).forEach((gainNode) => {
            gainNode.disconnect();
          });
          gainNodesRef.current = {};
        }

        if (panNodesRef.current) {
          Object.values(panNodesRef.current).forEach((panNode) => {
            panNode.disconnect();
          });
          panNodesRef.current = {};
        }

        if (scheduledBuffersRef.current) {
          Object.values(scheduledBuffersRef.current).forEach((buffer) => {
            buffer.forEach((entry) => {
              entry.sourceNode.stop();
              entry.sourceNode.disconnect();
            });
          });
          scheduledBuffersRef.current = {};
        }

        tearDownRef.current = () => {};
      };

      if (lastSeekFromTime.current !== null && isRealtimePlayback) {
        logSeekPlaybackChanges(lastSeekFromTime.current);
        lastSeekFromTime.current = null;
      }

      return { masterGainNode, fadeGainNode };
    },
    [
      state,
      volume,
      getCurrentTime,
      logSeekPlaybackChanges,
      logSongEndPlaybackChanges,
    ]
  );

  const stop = useCallback(
    (andSeekToStart: boolean = false) => {
      if (andSeekToStart) {
        lastStartedFromRef.current = 0;
      } else {
        lastStartedFromRef.current = getCurrentTime();
      }

      if (tearDownRef.current) {
        tearDownRef.current();
      }

      setPlaying(false);

      if (andSeekToStart) {
        lastSeekTimeRef.current = 0;
      }
    },
    [getCurrentTime]
  );

  const playDebounced = useDebounceCallback(play, 20);

  const seek = useCallback(
    (time: number, andUpdateWhilePlaying: boolean = true) => {
      lastSeekFromTime.current = getCurrentTime();
      if (time > state.skipTimeStart && time < state.skipTimeEnd) {
        time = Math.max(
          0,
          addBeatsToSeconds(-2, state.skipTimeStart, state.timing)
        );
      } else if (time > state.skipTimeEnd) {
        time -= state.skipTimeEnd - state.skipTimeStart;
      }
      lastSeekTimeRef.current = time;
      lastStartedFromRef.current = time;
      if (andUpdateWhilePlaying || !playingRef.current) {
        lastReportedTimeRef.current = time;
        lastReportedAtRef.current = Date.now();
      }
      if (andUpdateWhilePlaying && playingRef.current) {
        playDebounced(time);
      }
    },
    [play, state.skipTimeStart, state.skipTimeEnd]
  );

  useDismount(stop);

  useEffect(() => {
    audioBuffersByPojoBufferIdRef.current = {};
  }, [state.tracks]);

  const lastAmplitudeAutomationRef = useRef<AutomationPoint[]>(
    state.amplitudeAutomation
  );
  useEffect(() => {
    if (
      playingRef.current &&
      !isEqual(state.amplitudeAutomation, lastAmplitudeAutomationRef.current)
    ) {
      playDebounced(undefined, currentEndTimeRef.current);
    }
    lastAmplitudeAutomationRef.current = state.amplitudeAutomation;
  }, [state]);

  useEffect(() => {
    if (playingRef.current) {
      stop();
    }
  }, [state.skipTimeEnd, state.skipTimeStart]);

  const render = useCallback(async () => {
    // Create offline context with enough duration to render the full project
    const offlineCtx = new OfflineAudioContext(
      2, // stereo output
      Math.ceil(state.songEndSeconds * audioContext.sampleRate), // total samples needed
      audioContext.sampleRate
    );

    // Play the entire project from start to finish in the offline context
    const playResult = await play(0, state.songEndSeconds, offlineCtx);

    if (!playResult) {
      throw new Error('Failed to play');
    }

    const { masterGainNode, fadeGainNode } = playResult;

    // Start the rendering process
    const renderedBuffer = await offlineCtx.startRendering();

    // Clean up nodes
    masterGainNode.disconnect();
    fadeGainNode.disconnect();

    return renderedBuffer;
  }, [state.songEndSeconds, play, state.tracks]);

  return useMemo(
    () => ({
      play,
      seek,
      stop,
      playing,
      getCurrentTime,
      setVolume,
      volume,
      setTrackAmplitude,
      setTrackBalance,
      setTrackClips,
      deleteTrack,
      addTrack,
      render,
    }),
    [
      play,
      stop,
      seek,
      playing,
      getCurrentTime,
      volume,
      setVolume,
      setTrackAmplitude,
      setTrackBalance,
      setTrackClips,
      deleteTrack,
      addTrack,
      render,
    ]
  );
};

export const trackEffectivelyMuted = (
  track: EditTrack,
  anyTrackSolo: boolean
) => {
  if (track.solo) return false;
  if (track.mute) return true;
  if (anyTrackSolo) return true;
  return false;
};

export const getEffectiveAmplitude = (
  track: EditTrack,
  anyTrackSolo: boolean
) => {
  if (trackEffectivelyMuted(track, anyTrackSolo)) return 0;
  return track.amplitude;
};

export const useEditPlaybackContext = () => {
  const previewClipContext = useContext(PreviewClipContext);
  const selectionContext = useContext(SelectionContext);
  const stemsContext = useContext(StemsContext);
  const editModeContext = useContext(EditModeContext);
  const timing = useContext(EditTimingContext);

  const editorPlaybackState = useMemo<EditorPlaybackState>(() => {
    const anyTrackSolo =
      stemsContext.stemTracks.some((t) => t.solo) ||
      stemsContext.fullSongTrack.solo;

    const tracks: EditorPlaybackTrack[] = [
      {
        id: stemsContext.fullSongTrack.id,
        amplitude:
          editModeContext.editMode === 'stems'
            ? getEffectiveAmplitude(stemsContext.fullSongTrack, anyTrackSolo)
            : 1,
        balance:
          editModeContext.editMode === 'stems'
            ? stemsContext.fullSongTrack.balance
            : 0,
        clips:
          editModeContext.editMode === 'stems'
            ? stemsContext.fullSongTrack.arrangedClips
            : previewClipContext.arrangedClips,
      },
    ];

    if (editModeContext.editMode === 'stems') {
      stemsContext.stemTracks.forEach((stemTrack) => {
        tracks.push({
          id: stemTrack.id,
          amplitude: getEffectiveAmplitude(stemTrack, anyTrackSolo),
          balance: stemTrack.balance,
          clips: stemTrack.arrangedClips,
        });
      });
    }

    const lastClipEndSeconds = Math.max(
      ...tracks.map((t) => Math.max(...t.clips.map((c) => c.endSeconds)))
    );

    return {
      tracks,
      skipTimeStart:
        editModeContext.editMode === 'remove'
          ? selectionContext.selectionStartSeconds
          : 0,
      skipTimeEnd:
        editModeContext.editMode === 'remove'
          ? selectionContext.selectionEndSeconds
          : 0,
      songEndSeconds: lastClipEndSeconds,
      playbackEndSeconds:
        editModeContext.editMode === 'crop'
          ? selectionContext.selectionEndSeconds
          : lastClipEndSeconds,
      amplitudeAutomation:
        editModeContext.editMode === 'fadeOut'
          ? [
              {
                beats: 0,
                value: 1,
                curve: 0,
              },
              {
                beats: getBeatsFromZero(
                  selectionContext.selectionStartSeconds,
                  timing
                ),
                value: 1,
                curve: 1,
              },
              {
                beats: getBeatsFromZero(
                  previewClipContext.songEndSeconds,
                  timing
                ),
                value: 0,
                curve: 1,
              },
            ]
          : [
              {
                beats: 0,
                value: 1,
                curve: 0,
              },
            ],
      timing,
    };
  }, [
    timing,
    stemsContext.stemTracks,
    stemsContext.fullSongTrack,
    previewClipContext.arrangedClips,
    selectionContext.selectionStartSeconds,
    selectionContext.selectionEndSeconds,
    previewClipContext.songEndSeconds,
    editModeContext.editMode,
  ]);

  const player = useEditorStatePlayer(editorPlaybackState);

  useEffect(() => {
    const onKeyDown = (e: KeyboardEvent) => {
      if (
        (e.target as HTMLElement)?.matches(
          'input, textarea, [contenteditable]'
        ) ||
        document.activeElement instanceof HTMLTextAreaElement ||
        document.activeElement instanceof HTMLInputElement
      )
        return;
      if (e.key === ' ') {
        if (player.playing) {
          player.stop();
        } else {
          player.play();
        }
        e.preventDefault();
      }
    };
    window.addEventListener('keydown', onKeyDown);
    return () => window.removeEventListener('keydown', onKeyDown);
  }, [player]);

  useEffect(() => {
    if (
      !player.playing &&
      getEditCategory(editModeContext.editMode) !== 'no-selection'
    ) {
      let targetTime = selectionContext.selectionStartSeconds;
      if (editModeContext.editMode === 'remove') {
        targetTime = addBeatsToSeconds(-2, targetTime, timing);
      }
      player.seek(Math.max(0, targetTime), false);
    }
  }, [player.playing, selectionContext.selectionStartSeconds]);

  useEffect(() => {
    const previewArrangement = previewClipContext.arrangedClips.find(
      (clip) => clip.clipId === previewClipContext.previewClipId
    );
    if (previewArrangement) {
      player.play(previewArrangement.startSeconds);
    } else {
      player.stop();
    }
  }, [previewClipContext.arrangedClips, previewClipContext.previewClipId]);

  const { playbar } = useStores();

  useEffect(() => {
    player.setVolume(playbar.volume);
  }, [playbar.volume]);

  const lastStemTrackIDsRef = useRef<string[]>([]);

  useEffect(() => {
    if (editModeContext.editMode !== 'stems') {
      player.stop();
    } else {
      const anyTrackSolo =
        stemsContext.stemTracks.some((t) => t.solo) ||
        stemsContext.fullSongTrack.solo;

      stemsContext.stemTracks.forEach((track) => {
        if (!lastStemTrackIDsRef.current.includes(track.id)) {
          player.addTrack(track.id);
        }

        player.setTrackAmplitude(
          track.id,
          getEffectiveAmplitude(track, anyTrackSolo)
        );

        player.setTrackBalance(track.id, track.balance);

        player.setTrackClips(track.id, track.arrangedClips);
      });

      player.setTrackAmplitude(
        stemsContext.fullSongTrack.id,
        getEffectiveAmplitude(stemsContext.fullSongTrack, anyTrackSolo)
      );

      player.setTrackBalance(
        stemsContext.fullSongTrack.id,
        stemsContext.fullSongTrack.balance
      );

      player.setTrackClips(
        stemsContext.fullSongTrack.id,
        stemsContext.fullSongTrack.arrangedClips
      );
    }
    const nextIDs = stemsContext.stemTracks.map((t) => t.id);
    lastStemTrackIDsRef.current.forEach((id) => {
      if (!nextIDs.includes(id)) {
        player.deleteTrack(id);
      }
    });
    lastStemTrackIDsRef.current = nextIDs;
  }, [
    stemsContext.stemTracks,
    stemsContext.fullSongTrack,
    editModeContext.editMode,
  ]);

  return player;
};

const EditPlaybackContext = createContext<
  ReturnType<typeof useEditPlaybackContext>
>(undefined as never);

export const EditPlaybackContextProvider = ({
  children,
}: {
  children: React.ReactNode[] | React.ReactNode;
}) => {
  const value = useEditPlaybackContext();

  return (
    <EditPlaybackContext.Provider value={value}>
      {children}
    </EditPlaybackContext.Provider>
  );
};

export default EditPlaybackContext;