import deriveTiming from '@suno/studiokit/deriveTiming';
import { getSecondsFromZero } from '@suno/studiokit/timeMapping';
import { getWarpEnabledAndPopulated } from '@suno/studiokit/warpUtils';
import { uniq } from 'lodash-es';
import { createSelector, createSelectorCreator, lruMemoize } from 'reselect';

import { Clip } from '@/state/clipStore';

import { AlignedLyric } from '../edit2025/types';
import cachedGetLyricsBetween from './cachedGetLyricsBetween';
import { minimumTrackHeight } from './createTrack';
import reduceClips, { areClipsContentAligned } from './reduceClips';
import shouldShowTakeLanes from './shouldShowTakeLanes';
import {
  AFTER_LAST_TRACK,
  ClipArrangementPackage,
  StudioClip,
  StudioProjectState,
  StudioTrack,
  StudioTrackCore,
} from './types';
import { getClipAlignmentSpecFromHistory } from './useAlignedClips';

const shallowEqualArray = (a: any[], b: any[]) => {
  if (a.length !== b.length) return false;
  for (let i = 0; i < a.length; i++) {
    if (a[i] !== b[i]) return false;
  }
  return true;
};

const shallowEqualObject = (
  a: Record<string, any> | null | undefined,
  b: Record<string, any> | null | undefined
) => {
  if (a instanceof Object && b instanceof Object) {
    if (Object.keys(a).length !== Object.keys(b).length) return false;
    for (const key in a) {
      if (a[key] !== b[key]) return false;
    }
    return true;
  } else {
    return a == b;
  }
};

const createObjectSelector = createSelectorCreator({
  memoize: lruMemoize,
  memoizeOptions: {
    resultEqualityCheck: shallowEqualObject,
  },
});

const createArraySelector = createSelectorCreator({
  memoize: lruMemoize,
  memoizeOptions: {
    resultEqualityCheck: shallowEqualArray,
  },
});

export const FALLBACK_SONG_START_BEATS = 0;
export const FALLBACK_SONG_END_BEATS = 240;

const getState = (state: StudioProjectState) => state;

export const getSelection = createSelector(
  getState,
  (state) => state.selection,
  {
    memoizeOptions: {
      resultEqualityCheck: shallowEqualObject,
    },
  }
);

export const getFocusedTrackId = createSelector(
  getSelection,
  (selection) => selection.focusedTrackId
);

export const getSortedSelectionBounds = createArraySelector(
  getSelection,
  (selection) =>
    [selection.anchorBeats, selection.focusBeats].sort((a, b) => a - b)
);

export const getSelectionStartBeats = createSelector(
  getSortedSelectionBounds,
  (bounds) => bounds[0]
);

export const getSelectionEndBeats = createSelector(
  getSortedSelectionBounds,
  (bounds) => bounds[1]
);

export const getSelectionDurationBeats = createSelector(
  getSelectionEndBeats,
  getSelectionStartBeats,
  (selectionEndBeats, selectionStartBeats) =>
    selectionEndBeats - selectionStartBeats
);

export const getTiming = createObjectSelector(
  getState,
  (state) => state.timing
);

export const getTimingType = createSelector(getTiming, (timing) => timing.type);

export const getTracks = createArraySelector(getState, (state) => state.tracks);

export const getTracksAndTakeLanes = createArraySelector(getState, (state) => {
  const result: StudioTrackCore[] = [];
  for (const track of state.tracks) {
    result.push(track);
    result.push(...track.takeLanes);
  }
  return result;
});

export const getTracksAndExpandedTakeLanes = createArraySelector(
  getState,
  (state) => {
    const result: StudioTrackCore[] = [];
    for (const track of state.tracks) {
      result.push(track);
      if (track.takeLanesExpanded && track.height > minimumTrackHeight) {
        result.push(...track.takeLanes);
      }
    }
    return result;
  }
);

export const getTakeLaneParentsByTrackId = createObjectSelector(
  getTracks,
  (tracks) => {
    const result: Record<string, StudioTrack> = {};
    for (const track of tracks) {
      result[track.id] = track;
      for (const takeLane of track.takeLanes) {
        result[takeLane.id] = track;
      }
    }
    return result;
  }
);

export const getSelectionConnectedTrackAndTakelaneIds = createSelector(
  [getTracks, getSelection],
  (tracks, selection) => {
    const result: string[] = [];
    const selectedParentTracks = tracks.filter(
      (t) =>
        selection.trackIds.includes(t.id) ||
        t.takeLanes.some((tl) => selection.trackIds.includes(tl.id))
    );
    for (const track of selectedParentTracks) {
      result.push(track.id);
      for (const takeLane of track.takeLanes) {
        result.push(takeLane.id);
      }
    }
    return result;
  }
);

const invertTimespans = (
  timespans: { startBeats: number; endBeats: number }[]
) => {
  if (timespans.length === 0) {
    return [{ startBeats: -Infinity, endBeats: Infinity }];
  }

  // Sort timespans by startBeats
  const sortedTimespans = [...timespans].sort(
    (a, b) => a.startBeats - b.startBeats
  );

  const inverted: { startBeats: number; endBeats: number }[] = [];

  // Add gap from -Infinity to first timespan
  if (sortedTimespans[0].startBeats > -Infinity) {
    inverted.push({
      startBeats: -Infinity,
      endBeats: sortedTimespans[0].startBeats,
    });
  }

  // Add gaps between consecutive timespans
  for (let i = 0; i < sortedTimespans.length - 1; i++) {
    const current = sortedTimespans[i];
    const next = sortedTimespans[i + 1];
    if (current.endBeats < next.startBeats) {
      inverted.push({
        startBeats: current.endBeats,
        endBeats: next.startBeats,
      });
    }
  }

  // Add gap from last timespan to Infinity
  const lastTimespan = sortedTimespans[sortedTimespans.length - 1];
  if (lastTimespan.endBeats < Infinity) {
    inverted.push({
      startBeats: lastTimespan.endBeats,
      endBeats: Infinity,
    });
  }

  return inverted;
};

export const getUnarrangedTakeTimespansByTakeLaneId = createObjectSelector(
  getTracks,
  (tracks) => {
    const result: Record<string, { startBeats: number; endBeats: number }[]> =
      {};
    for (const track of tracks) {
      if (shouldShowTakeLanes(track)) {
        for (const takeLane of track.takeLanes) {
          const arrangedTimespans: { startBeats: number; endBeats: number }[] =
            [];

          takeLane.clips.forEach((clip) => {
            const overlappingTrackClips = track.clips.filter(
              (c) =>
                c.startBeats < clip.endBeats && c.endBeats > clip.startBeats
            );
            overlappingTrackClips.forEach((c) => {
              if (areClipsContentAligned(c, clip)) {
                arrangedTimespans.push({
                  startBeats: Math.max(c.startBeats, clip.startBeats),
                  endBeats: Math.min(c.endBeats, clip.endBeats),
                });
              }
            });
          });

          result[takeLane.id] = invertTimespans(arrangedTimespans);
        }
      }
    }
    return result;
  }
);

export const getUnfocusedTrackIds = createSelector(getSelection, (selection) =>
  selection.trackIds.filter((id) => id !== selection.focusedTrackId)
);

export const getUnfocusedTracks = createSelector(
  [getTracks, getFocusedTrackId],
  (tracks, focusedTrackId) => tracks.filter((t) => t.id !== focusedTrackId)
);

export const getCanSelectFollowTrackTiming = createSelector(
  getTracks,
  (tracks) => tracks.length > 0
);

export const getDerivedTiming = createObjectSelector(
  [getTracks, getTiming],
  (tracks, timing) => deriveTiming({ timing, tracks })
);

export const getFallbackBPS = createSelector(
  getDerivedTiming,
  (derivedTiming) => derivedTiming.bps
);

export const getSelectionStartSeconds = createSelector(
  [getSelectionStartBeats, getDerivedTiming],
  (selectionStartBeats, timing) =>
    getSecondsFromZero(selectionStartBeats, timing)
);

export const getSelectionEndSeconds = createSelector(
  [getSelectionEndBeats, getDerivedTiming],
  (selectionEndBeats, timing) => getSecondsFromZero(selectionEndBeats, timing)
);

export const getSelectionDurationSeconds = createSelector(
  [getSelectionEndSeconds, getSelectionStartSeconds],
  (selectionEndSeconds, selectionStartSeconds) =>
    selectionEndSeconds - selectionStartSeconds
);

export const getAllUsedClipIds = createArraySelector(
  getTracksAndTakeLanes,
  (tracks) =>
    uniq(
      tracks.flatMap((t) =>
        t.clips
          .map((c) => c.clipId)
          .concat(t.clipCreationIntents.flatMap((c) => c.possibleClipIds))
      )
    ).filter(Boolean) as string[]
);

export const getTracksById = createObjectSelector(getTracks, (tracks) =>
  Object.fromEntries(tracks.map((t) => [t.id, t]))
);

export const getTakeLanesById = createObjectSelector(getTracks, (tracks) =>
  Object.fromEntries(
    tracks.flatMap((t) => t.takeLanes.map((tl) => [tl.id, tl]))
  )
);

export const getTracksAndTakeLanesById = createObjectSelector(
  getTracksAndTakeLanes,
  (tracksAndTakeLanes) =>
    Object.fromEntries(tracksAndTakeLanes.map((t) => [t.id, t]))
);

export const getSelectionTrackIds = createArraySelector(
  [getSelection, getTracksAndTakeLanesById],
  (selection, tracksAndTakeLanesById) =>
    selection.trackIds.filter((t) => !!tracksAndTakeLanesById[t])
);

export const getTrackIds = createArraySelector(getTracks, (tracks) =>
  tracks.map((t) => t.id)
);

export const getTrackAndExpandedTakeLaneIds = createArraySelector(
  getTracksAndExpandedTakeLanes,
  (tracksAndExpandedTakeLanes) => tracksAndExpandedTakeLanes.map((t) => t.id)
);

export const getTrackAndTakeLaneIds = createArraySelector(
  getTracksAndTakeLanes,
  (tracksAndTakeLanes) => tracksAndTakeLanes.map((t) => t.id)
);

export const getSelectableTrackIds = createArraySelector(
  getTrackAndExpandedTakeLaneIds,
  (trackIds) => trackIds.concat(AFTER_LAST_TRACK)
);

export const getFocusedTrack = createSelector(
  [getFocusedTrackId, getTracksById],
  (trackId, tracksById) => (trackId ? (tracksById[trackId] ?? null) : null)
);

export const getAnyTrackSolo = createSelector(getTracks, (tracks) =>
  tracks.some((t) => t.solo)
);

export const getEffectivelyMutedTracks = createObjectSelector(
  [getTracks, getAnyTrackSolo],
  (tracks, anyTrackSolo) =>
    Object.fromEntries(
      tracks.map((t) => [t.id, anyTrackSolo ? !t.solo : t.mute])
    )
);

export const getStudioClips = createArraySelector(
  getTracksAndTakeLanes,
  (tracks) => tracks.flatMap((t) => t.clips)
);

export const getUnwarpedStudioClips = createArraySelector(
  getStudioClips,
  (studioClips) =>
    studioClips.filter((c) => !getWarpEnabledAndPopulated(c.warp))
);

export const getStudioClipsById = createObjectSelector(
  getStudioClips,
  (studioClips) => Object.fromEntries(studioClips.map((c) => [c.id, c]))
);

export const getStudioClipIds = createArraySelector(
  getStudioClipsById,
  (studioClipsById) => Object.keys(studioClipsById)
);

export const getStudioClipsByClipId = createObjectSelector(
  getStudioClips,
  (studioClips) => {
    const result: Record<string, StudioClip[]> = {};
    for (const clip of studioClips) {
      if (clip.clipId) {
        (result[clip.clipId] ||= []).push(clip);
      }
    }
    return result;
  }
);

export const getStudioClipsByTrackId = createObjectSelector(
  getTracksAndTakeLanes,
  (tracks) => Object.fromEntries(tracks.map((t) => [t.id, t.clips]))
);

export const getStudioClipsByTrackOrTakeLaneId = createObjectSelector(
  getTracksAndTakeLanes,
  (tracks) => Object.fromEntries(tracks.map((t) => [t.id, t.clips]))
);

export const getClipIdsByStudioClipId = createObjectSelector(
  getStudioClipsById,
  (studioClipsById) =>
    Object.fromEntries(
      Object.entries(studioClipsById).map(([studioClipId, studioClip]) => [
        studioClipId,
        studioClip.clipId,
      ])
    )
);

export const getSortedClipsByTrackId = createObjectSelector(
  getStudioClipsByTrackId,
  (studioClipsByTrackId) =>
    Object.fromEntries(
      Object.entries(studioClipsByTrackId).map(([trackId, clips]) => [
        trackId,
        clips.sort((a, b) => a.startBeats - b.startBeats),
      ])
    )
);

export const getClipCreationIntentsByTrackId = createObjectSelector(
  getTracks,
  (tracks) =>
    Object.fromEntries(tracks.map((t) => [t.id, t.clipCreationIntents]))
);

export const getClipCreationIntentsByTrackOrTakeLaneId = createObjectSelector(
  getTracksAndTakeLanes,
  (tracks) =>
    Object.fromEntries(tracks.map((t) => [t.id, t.clipCreationIntents]))
);

export const getSelectedClips = createArraySelector(
  [
    getSelectionTrackIds,
    getSelectionStartBeats,
    getSelectionEndBeats,
    getStudioClipsByTrackId,
  ],
  (trackIds, selectionStartBeats, selectionEndBeats, studioClipsByTrackId) =>
    trackIds.flatMap((id) =>
      (studioClipsByTrackId[id] || [])
        .filter(
          (c) =>
            c.startBeats < selectionEndBeats && c.endBeats > selectionStartBeats
        )
        .sort((a, b) => a.startBeats - b.startBeats)
    )
);

export const getFocusedTrackClips = createArraySelector(
  [getFocusedTrackId, getStudioClipsByTrackId],
  (focusedTrackId, studioClipsByTrackId) =>
    focusedTrackId ? studioClipsByTrackId[focusedTrackId] || [] : []
);

export const getSelectedClipCreationIntents = createArraySelector(
  [
    getSelectionTrackIds,
    getSelectionStartBeats,
    getSelectionEndBeats,
    getClipCreationIntentsByTrackOrTakeLaneId,
  ],
  (
    trackIds,
    selectionStartBeats,
    selectionEndBeats,
    clipCreationIntentsByTrackId
  ) =>
    trackIds.flatMap((id) =>
      (clipCreationIntentsByTrackId[id] || []).filter(
        (c) =>
          (c.startBeats ?? -Infinity) < selectionEndBeats &&
          (c.endBeats ?? Infinity) > selectionStartBeats
      )
    )
);

export const getEntirelySelectedClips = createArraySelector(
  [getSelectionStartBeats, getSelectionEndBeats, getSelectedClips],
  (selectionStartBeats, selectionEndBeats, clips) =>
    clips.filter(
      (c) =>
        c.startBeats >= selectionStartBeats && c.endBeats <= selectionEndBeats
    )
);

export const getSelectedTracks = createArraySelector(
  [getSelectionTrackIds, getTracksById],
  (trackIds, tracksById) => trackIds.map((id) => tracksById[id]).filter(Boolean)
);

export const getFocusedStudioClip = createSelector(
  [getSelectionStartBeats, getSelectionEndBeats, getSelectedClips],
  (selectionStartBeats, selectionEndBeats, clips) =>
    clips.find(
      (c) =>
        c.startBeats === selectionStartBeats && c.endBeats === selectionEndBeats
    ) || null
);

export const getFocusedStudioClipClipId = createSelector(
  [getFocusedStudioClip],
  (focusedStudioClip) => focusedStudioClip?.clipId || null
);

export const getEarliestClipBeats = createSelector(
  getStudioClips,
  (studioClips) =>
    studioClips.reduce<number>(
      (acc, c) => Math.min(acc, c.startBeats),
      Infinity
    )
);

export const getLatestClipBeats = createSelector(
  getStudioClips,
  (studioClips) =>
    studioClips.reduce<number>((acc, c) => Math.max(acc, c.endBeats), -Infinity)
);

export const getSongStartBeats = createSelector(
  getEarliestClipBeats,
  (earliestClipBeats) =>
    Number.isFinite(earliestClipBeats)
      ? earliestClipBeats
      : FALLBACK_SONG_START_BEATS
);

export const getSongEndBeats = createSelector(
  getLatestClipBeats,
  (latestClipBeats) =>
    Number.isFinite(latestClipBeats) ? latestClipBeats : FALLBACK_SONG_END_BEATS
);

export const getSongStartSeconds = createSelector(
  [getSongStartBeats, getDerivedTiming],
  (songStartBeats, timing) => getSecondsFromZero(songStartBeats, timing)
);

export const getSongEndSeconds = createSelector(
  [getSongEndBeats, getDerivedTiming],
  (songEndBeats, timing) => getSecondsFromZero(songEndBeats, timing)
);

export const getStudioDisplayedEndBeats = createSelector(
  [getSongEndBeats, getLatestClipBeats],
  (songEndBeats, latestClipBeats) =>
    Math.max(
      Number.isFinite(latestClipBeats) ? songEndBeats * 1.5 : songEndBeats,
      FALLBACK_SONG_END_BEATS
    )
);

export const getStudioDisplayedStartBeats = createSelector(
  [getSongStartBeats],
  (songStartBeats) => Math.min(songStartBeats, 0)
);

export const getStudioDisplayedEndSeconds = createSelector(
  [getStudioDisplayedEndBeats, getDerivedTiming],
  (studioDisplayedEndBeats, timing) =>
    getSecondsFromZero(studioDisplayedEndBeats, timing)
);

export const getStudioDisplayedStartSeconds = createSelector(
  [getStudioDisplayedStartBeats, getDerivedTiming],
  (studioDisplayedStartBeats, timing) =>
    getSecondsFromZero(studioDisplayedStartBeats, timing)
);

export const getFocusedClipCreationIntent = createSelector(
  [
    getSelectionStartBeats,
    getSelectionEndBeats,
    getSelectedClipCreationIntents,
    getSongStartBeats,
    getSongEndBeats,
  ],
  (
    selectionStartBeats,
    selectionEndBeats,
    clipCreationIntents,
    songStartBeats,
    songEndBeats
  ) =>
    clipCreationIntents.find((c) => {
      let startMatches = true;
      let selectionEndBeatsMatches = true;
      if (c.startBeats !== undefined && c.startBeats !== selectionStartBeats) {
        startMatches = false;
      } else if (
        c.startBeats === undefined &&
        selectionStartBeats > songStartBeats
      ) {
        startMatches = false;
      }

      if (c.endBeats !== undefined && c.endBeats !== selectionEndBeats) {
        selectionEndBeatsMatches = false;
      } else if (c.endBeats === undefined && selectionEndBeats < songEndBeats) {
        selectionEndBeatsMatches = false;
      }

      return startMatches && selectionEndBeatsMatches;
    }) || null
);

export const getLastClipBeforeSelection = createSelector(
  [getSelectionStartBeats, getFocusedTrackClips],
  (selectionStartBeats, clips) =>
    clips.findLast((c) => c.endBeats <= selectionStartBeats) || null
);

export const getFirstClipAfterSelection = createSelector(
  [getSelectionEndBeats, getFocusedTrackClips],
  (selectionEndBeats, clips) =>
    clips.find((c) => c.startBeats >= selectionEndBeats) || null
);

export const getTimingConstant = createSelector([getTiming], (timing) =>
  timing.type === 'manual' ? Boolean(timing.lockBPS) : false
);

export const getEarliestUsedBeat = createSelector(
  [getEarliestClipBeats, getTiming],
  (earliestClipBeats, timing) => {
    const value = Math.min(
      earliestClipBeats,
      ...(timing.type === 'manual'
        ? timing.bpsAutomation.map((b) => b.beats)
        : [])
    );
    return Number.isFinite(value) ? value : 0;
  }
);

export const getSelectionGapStartBeats = createSelector(
  getLastClipBeforeSelection,
  (lastClipBeforeSelection) =>
    lastClipBeforeSelection ? lastClipBeforeSelection.endBeats : null
);

export const getSelectionGapEndBeats = createSelector(
  getFirstClipAfterSelection,
  (firstClipAfterSelection) =>
    firstClipAfterSelection ? firstClipAfterSelection.startBeats : null
);

export const getSelectionGapSize = createSelector(
  [getSelectionGapStartBeats, getSelectionGapEndBeats],
  (selectionGapStartBeats, selectionGapEndBeats) =>
    selectionGapEndBeats !== null && selectionGapStartBeats !== null
      ? selectionGapEndBeats - selectionGapStartBeats
      : 0
);

export const getReducedStudioClipsByTrackId = createObjectSelector(
  [getTrackIds, getSortedClipsByTrackId],
  (trackIds, sortedClipsByTrackId) =>
    Object.fromEntries(
      trackIds.map((trackId) => [
        trackId,
        reduceClips(sortedClipsByTrackId[trackId] || []),
      ])
    )
);

const getStateAndClipsById = (arg: {
  state: StudioProjectState;
  clipsById: Record<string, Clip>;
}) => arg;

const getCachedStateAndClipsById = createObjectSelector(
  getStateAndClipsById,
  (x) => x
);

export const getFocusedClip = createSelector(
  getCachedStateAndClipsById,
  ({ state, clipsById }) => {
    const focusedStudioClipClipId = getFocusedStudioClipClipId(state);
    return focusedStudioClipClipId
      ? clipsById[focusedStudioClipClipId] || null
      : null;
  }
);

const getStateAndAlignedLyricsByClipId = (arg: {
  state: StudioProjectState;
  alignedLyricsByClipId: Record<string, AlignedLyric[]>;
}) => arg;

const getCachedStateAndAlignedLyricsByClipId = createObjectSelector(
  getStateAndAlignedLyricsByClipId,
  (x) => x
);

export const getAlignedLyricsByTrackId = createObjectSelector(
  getCachedStateAndAlignedLyricsByClipId,
  ({ state, alignedLyricsByClipId }) => {
    const reducedStudioClipsByTrackId = getReducedStudioClipsByTrackId(state);
    const timing = getDerivedTiming(state);
    return Object.fromEntries(
      state.tracks.map((t) => [
        t.id,
        cachedGetLyricsBetween(
          timing,
          reducedStudioClipsByTrackId[t.id] || [],
          alignedLyricsByClipId,
          -Infinity,
          Infinity
        ),
      ])
    );
  }
);

export const getFocusedTrackAlignedLyrics = createSelector(
  [
    getCachedStateAndAlignedLyricsByClipId,
    ({ state }) => getDerivedTiming(state),
    ({ state }) => {
      const focusedTrackId = getFocusedTrackId(state);
      return focusedTrackId
        ? getReducedStudioClipsByTrackId(state)[focusedTrackId] || []
        : [];
    },
  ],
  ({ alignedLyricsByClipId }, timing, reducedStudioClips) => {
    return cachedGetLyricsBetween(
      timing,
      reducedStudioClips,
      alignedLyricsByClipId,
      -Infinity,
      Infinity
    );
  }
);

export const getCachedClipAlignmentSpecFromHistory = createObjectSelector(
  getFocusedClip,
  (clip) => (clip ? getClipAlignmentSpecFromHistory(clip) : undefined)
);

export const getFocusedArrangementPackage = createObjectSelector(
  [
    getFocusedClip,
    getCachedClipAlignmentSpecFromHistory,
    ({ state }) => getFocusedStudioClip(state),
  ],
  (clip, alignmentSpec, studioClip) =>
    clip && studioClip
      ? ({
          clip,
          studioClip,
          alignmentSpec,
        } as ClipArrangementPackage)
      : null
);

export const getClipIdsReadyForDownbeatAnalysis = createArraySelector(
  [getCachedStateAndClipsById, ({ state }) => getStudioClips(state)],
  ({ clipsById }, studioClips) =>
    uniq(
      studioClips
        .filter((c) => c.warp.awaitingAnalysis && c.clipId)
        .filter((c) =>
          ['complete', 'streaming'].includes(clipsById[c.clipId!]?.status ?? '')
        )
        .map((c) => c.clipId) as string[]
    )
);

export const getStreamingClipIds = createArraySelector(
  getStudioClips,
  (studioClips) =>
    studioClips
      .filter((c) => c.streaming && c.clipId)
      .map((c) => c.clipId) as string[]
);
const getTrack = (track: StudioTrack) => track;

export const getHasSoloedTakeLane = createSelector(
  [getTrack],
  (track) =>
    track.soloTakeLaneId &&
    track.takeLanes.some((tl) => tl.id === track.soloTakeLaneId) &&
    shouldShowTakeLanes(track)
);

export const getTotalTrackHeight = createSelector(
  [getTracksAndExpandedTakeLanes],
  (tracksAndExpandedTakeLanes) =>
    tracksAndExpandedTakeLanes.reduce((acc, track) => acc + track.height, 0)
);