import styled from '@emotion/styled';
import { getBeatsFromZero } from '@suno/studiokit/timeMapping';
import { observer } from 'mobx-react-lite';
import { Fragment, useCallback, useContext, useRef, useState } from 'react';
// eslint-disable-next-line @suno-custom/no-react-icons -- grandfathered in, need to migrate
import { BiFolder, BiSolidEraser } from 'react-icons/bi';

import { useStores } from '@/app/(root)/AppProviders';
import Button, { ButtonShape, ButtonVariant } from '@/components/button/Button';
import {
  ContextMenuItem,
  ContextMenuTrigger,
} from '@/components/contextMenu/ContextMenu';
import {
  Divider,
  Panel,
  PanelSize,
  VerticalPanelGroup,
} from '@/components/panelSystem/PanelSystem';
import useClip from '@/hooks/useClip';
import useUploadAndInitializeClip from '@/hooks/useUploadAndInitializeClip';
import { FourStarIcon, ShareArrowIcon } from '@/icons';
import { CreateIcon, MoreHorizontalIcon } from '@/icons';
import makeWavBuffer from '@/lib/makeWavBuffer';
import { PlanFeature } from '@/state/sessionStore';
import { isFeatureEnabledForPlan } from '@/utils/session';

import SpinnerSVG from '../svg/SpinnerSVG';
import { Tooltip } from '../tooltip/Tooltip';
import AddTracksModal from './AddTracksModal';
import EditModeContext from './EditModeContext';
import EditPlaybackContext from './EditPlaybackContext';
import EditTimelineContext from './EditTimelineContext';
import MultiTrackTrackHeader from './MultiTrackTrackHeader';
import PreviewClipContext from './PreviewClipContext';
import SelectedTrackSidebar from './SelectedTrackSidebar';
import SelectionContext from './SelectionContext';
import StemsContext, { smallTrackHeight } from './StemsContext';
import TimelineContainer from './TimelineContainer';
import CanvasRenderer from './canvasRenderer/CanvasRenderer';
import { constructEditCanvasRenderContext } from './canvasRenderer/EditCanvasRenderContext';
import ScrollCatcher from './canvasRenderer/ScrollCatcher';
import applyAutoScrollLimiting from './canvasRenderer/applyAutoScrollLimiting';
import makeBaseClickHandlerRegion from './canvasRenderer/makeBaseClickHandlerRegion';
import makeBeatGridRegion from './canvasRenderer/makeBeatGridRegion';
import makePlayheadRegion from './canvasRenderer/makePlayheadRegion';
import makeScrollbarRegion from './canvasRenderer/makeScrollbarRegion';
import makeSectionRegions from './canvasRenderer/makeSectionRegions';
import makeStemRegions from './canvasRenderer/makeStemRegions';
import {
  CanvasRendererContainer,
  StyledTrackHeader,
  TrackHeaders,
} from './components';
import AlignedLyricsContext from './lyrics/AlignedLyricsContext';
import EditTimingContext from './queryHooks/EditTimingContext';

const MultiTrackToolbarWrapper = styled.div<{ fadeOut: boolean }>`
  opacity: ${({ fadeOut }) => (fadeOut ? 0 : 1)};
  transition: opacity 0.4s;
  display: flex;
  align-items: center;
  justify-content: space-between;
  gap: 10px;
  padding: 0 10px;
`;

const ToolbarSection = styled.div`
  display: flex;
  align-items: center;
  gap: 10px;
`;

const MultiTrackTimelineWrapper = styled.div`
  border-top: 1px solid rgba(255, 255, 255, 0.1);
  display: grid;
  grid-template-columns: 200px 1fr;
  min-height: 100%;
`;

const MultiTrackTrackSection = styled.div`
  background-color: #1a1a1a;
  display: grid;
  min-height: 0;
  grid-template-rows: 1fr;
  border-right: 1px solid rgba(255, 255, 255, 0.1);
`;

const MultiTrackTrackHeaders = styled(TrackHeaders)`
  background-color: transparent;
  border-right: none;
`;

const MultiTrackHeaderListButton = styled(
  StyledTrackHeader.withComponent('button')
)`
  border: none;
  cursor: pointer;
  width: 100%;
  color: inherit;
  background: none;
  justify-content: flex-start;
  text-align: left;
  font-size: 10px;
  &:hover {
    background-color: rgba(255, 255, 255, 0.1);
  }
`;

const TrackDivider = styled.div`
  height: 1px;
  width: 100%;
  background-color: #333;
`;

const SectionDivider = styled.div`
  height: 3px;
  width: 100%;
  position: relative;
  &:after {
    content: '';
    display: block;
    position: absolute;
    left: 0;
    top: 0;
    bottom: 0;
    width: 100vw;
    background-color: #333;
    z-index: 2;
  }
`;

const StemsList = styled.div`
  position: absolute;
  top: 0;
  bottom: 0;
  left: 0;
  right: 0;
  overflow-y: auto;
`;

const WalkthroughWrapper = styled.div`
  position: absolute;
  top: 0;
  left: 200px;
  right: 0;
  height: 320px;
  z-index: 2;
  display: grid;
  grid-template-columns: 1fr 1fr 1fr;
  gap: 5px;
  padding: 60px;
`;

const WalkthroughCard = styled.div`
  padding: 20px;
  border-radius: 10px;
  backdrop-filter: blur(20px);
  background-color: rgba(255, 255, 255, 0);
  transition: background-color 0.2s;
  border: 1px solid rgba(255, 255, 255, 0.05);
  background-image: radial-gradient(
    circle at 50% 100%,
    #4658a588 20%,
    #ffffff02 100%
  );
  &:hover {
    cursor: pointer;
    background-color: rgba(255, 255, 255, 0.05);
  }
`;

const WalkthroughTitle = styled.div`
  font-size: 28px;
  font-family: 'PP Editorial New', ui-serif, Georgia, Cambria;
  margin-top: 10px;
`;

const WalkthroughDescription = styled.div`
  font-size: 14px;
  opacity: 0.5;
`;

const smallPanelSize: PanelSize = {
  unit: 'px',
  count: smallTrackHeight,
};

const oneFRPanelSize: PanelSize = {
  unit: 'fr',
  count: 1,
};

export default observer(function MultiTrackTimeline({
  clipId,
  fadeOut,
}: {
  clipId: string;
  fadeOut: boolean;
}) {
  const editTimelineContext = useContext(EditTimelineContext);
  const {
    frameCountRef,
    scrollXRef,
    pxPerBeatRef,
    resetZoom,
    handleScroll,
    handleZoom,
    wrapperRef,
    zoomAboutSelectionOrCenter,
  } = editTimelineContext;

  const { edit, clips, session } = useStores();
  const previewClipContext = useContext(PreviewClipContext);
  const stemsContext = useContext(StemsContext);
  const { editMode } = useContext(EditModeContext);
  const alignedLyricsContext = useContext(AlignedLyricsContext);
  const timing = useContext(EditTimingContext);

  const playbackContext = useContext(EditPlaybackContext);
  const selectionContext = useContext(SelectionContext);
  const sectionDividerRef = useRef<HTMLDivElement>(null);

  const fullSongPanelSize: PanelSize = {
    unit: 'px',
    count: stemsContext.fullSongTrack.height,
  };

  const sectionHeaderRef = useRef<HTMLDivElement>(null);
  const waveformHeaderRef = useRef<HTMLDivElement>(null);
  const stemsListRef = useRef<HTMLDivElement>(null);

  const getCanvasRegions = useCallback(() => {
    const wrapper = wrapperRef.current;
    const waveformHeader = waveformHeaderRef.current;
    const sectionHeader = sectionHeaderRef.current;
    const stemsList = stemsListRef.current;
    const sectionDivider = sectionDividerRef.current;

    if (
      !wrapper ||
      !waveformHeader ||
      !sectionHeader ||
      !stemsList ||
      !sectionDivider
    )
      return [];

    const timelineWrapperRect = wrapper.getBoundingClientRect();

    const rectMinusSidebar = new DOMRect(
      timelineWrapperRect.left,
      timelineWrapperRect.top,
      timelineWrapperRect.width - (stemsContext.selectedTrackId ? 370 : 0),
      timelineWrapperRect.height
    );

    const songStartBeats = getBeatsFromZero(0, timing);
    const songEndBeats = getBeatsFromZero(
      previewClipContext.songEndSeconds,
      timing
    );

    applyAutoScrollLimiting(
      scrollXRef,
      pxPerBeatRef,
      rectMinusSidebar,
      songStartBeats,
      songEndBeats,
      zoomAboutSelectionOrCenter
    );

    const canvasRenderContext = constructEditCanvasRenderContext({
      editTimelineContext,
      selectionContext,
      previewClipContext,
      stemsContext,
      alignedLyricsContext,
      playbackContext,
      timing,

      timelineWrapperRect: wrapper.getBoundingClientRect(),
      timelineWrapper: wrapper,
      waveformHeader,
      sectionHeader,
      sectionDivider,
      stemRows: Array.from(stemsList.querySelectorAll('.stem-row')),
      clipId,
      editMode,
      editSessionId: edit.editSessionId,
    });

    return [
      makeBaseClickHandlerRegion(canvasRenderContext),
      makeBeatGridRegion(canvasRenderContext),
      ...makeSectionRegions(canvasRenderContext),
      ...makeStemRegions(canvasRenderContext),
      makePlayheadRegion(canvasRenderContext),
      makeScrollbarRegion(canvasRenderContext),
    ];
  }, [
    clipId,
    editMode,
    edit.editSessionId,
    editTimelineContext,
    selectionContext,
    previewClipContext,
    stemsContext,
    alignedLyricsContext,
    playbackContext,
    timing,
    zoomAboutSelectionOrCenter,
  ]);

  const receiveWrapperRef = useCallback(
    (el: HTMLDivElement | null) => {
      wrapperRef.current = el;
      if (el) {
        resetZoom();
      }
    },
    [resetZoom]
  );

  const [addTracksTask, setAddTracksTask] = useState<
    'add' | 'extract' | 'remove' | null
  >(null);

  const uploadAndInitializeClip = useUploadAndInitializeClip();

  const { clip } = useClip(clipId);
  const editModeContext = useContext(EditModeContext);

  const handleSaveToLibrary = useCallback(async () => {
    if (!clip?.title) return;
    const filename = `${clip.title}.wav`;
    const buffer = await playbackContext.render();
    if (!buffer) return;
    editModeContext.setLastAppliedEditType('stems');
    previewClipContext.setApplying(true);
    const wavBuffer = await makeWavBuffer(buffer);
    const file = new File([wavBuffer], filename, { type: 'audio/wav' });
    console.log('Got file, uploading...');
    const result = await uploadAndInitializeClip(
      file,
      'file_upload',
      undefined,
      `${clip.title} (Remix)`,
      clip.metadata.prompt || undefined,
      clipId,
      [stemsContext.fullSongTrack, ...stemsContext.stemTracks]
        .map((track) => track.arrangedClips.map((clip) => clip.clipId))
        .flat(2)
    );
    console.log('Upload complete, result:', result);
    window.location.href = `/song/${result.clipId}`;
  }, [clip, clips, playbackContext, uploadAndInitializeClip, editModeContext]);

  const receiveScroll = useCallback(
    (x: number, y: number) => {
      handleScroll(x, y);
      stemsListRef.current?.scrollBy(x, y);
    },
    [handleScroll]
  );

  const [downloadingZip, setDownloadingZip] = useState(false);

  if (!isFeatureEnabledForPlan(session, PlanFeature.EditMode)) {
    return (
      <div className='flex h-full items-center justify-center gap-2'>
        <Button href='/account' variant={ButtonVariant.Primary}>
          Subscribe
        </Button>
        to enable stems.
      </div>
    );
  }

  return (
    <TimelineContainer clipId={clipId}>
      {addTracksTask && (
        <AddTracksModal
          onClose={() => setAddTracksTask(null)}
          initialTask={addTracksTask}
        />
      )}

      <MultiTrackToolbarWrapper fadeOut={fadeOut}>
        <ToolbarSection />
        <ToolbarSection>
          <Button
            shape={ButtonShape.Pill}
            icon={CreateIcon}
            onClick={() => setAddTracksTask('add')}
          >
            Add Track
          </Button>
          <Button
            shape={ButtonShape.Pill}
            icon={ShareArrowIcon}
            onClick={() => setAddTracksTask('extract')}
          >
            Extract
          </Button>
          <Button
            shape={ButtonShape.Pill}
            icon={BiSolidEraser}
            onClick={() => setAddTracksTask('remove')}
          >
            Remove
          </Button>
        </ToolbarSection>
        <ToolbarSection>
          <ContextMenuTrigger
            title='Export Options'
            ButtonComponent={(props) => (
              <Button {...props}>
                <MoreHorizontalIcon />
              </Button>
            )}
            ContentsComponent={({ onClose }) => (
              <ContextMenuItem
                keepMenusOpen
                disabled={downloadingZip}
                icon={
                  downloadingZip ? (
                    <SpinnerSVG className='h-5 w-5' />
                  ) : (
                    <BiFolder className='h-5 w-5' />
                  )
                }
                onClick={async () => {
                  setDownloadingZip(true);
                  try {
                    await stemsContext.downloadZip();
                  } finally {
                    setDownloadingZip(false);
                    onClose();
                  }
                }}
              >
                Download stems as .zip
              </ContextMenuItem>
            )}
          />
          <Tooltip label='Apply all edits and mix changes into a new song in your library.'>
            <Button
              shape={ButtonShape.Pill}
              variant={ButtonVariant.Aura}
              icon={CreateIcon}
              onClick={() => handleSaveToLibrary()}
            >
              Export to Library
            </Button>
          </Tooltip>
        </ToolbarSection>
      </MultiTrackToolbarWrapper>
      <MultiTrackTimelineWrapper>
        <MultiTrackTrackSection>
          <MultiTrackTrackHeaders onWheel={() => frameCountRef.current++}>
            <VerticalPanelGroup frameCountRef={frameCountRef}>
              <Panel size={{ unit: 'px', count: 35 }}>
                <div />
              </Panel>
              <Divider influence='neither'>
                <TrackDivider />
              </Divider>
              <Panel size={smallPanelSize}>
                <StyledTrackHeader ref={sectionHeaderRef}>
                  Song Structure
                </StyledTrackHeader>
              </Panel>
              <Divider influence='neither'>
                <TrackDivider />
              </Divider>
              <Panel size={fullSongPanelSize}>
                <div className='h-full' ref={waveformHeaderRef}>
                  <MultiTrackTrackHeader
                    trackId={stemsContext.fullSongTrack.id}
                  />
                </div>
              </Panel>
              <Panel size={smallPanelSize}>
                <MultiTrackHeaderListButton
                  onClick={() => setAddTracksTask('remove')}
                >
                  - Remove Instruments
                </MultiTrackHeaderListButton>
              </Panel>
              <Divider influence='neither'>
                <SectionDivider ref={sectionDividerRef} />
              </Divider>
              <Panel size={oneFRPanelSize}>
                <div className='relative h-full'>
                  <StemsList ref={stemsListRef}>
                    <VerticalPanelGroup frameCountRef={frameCountRef}>
                      {stemsContext.stemTracks.map((track) => (
                        <Fragment key={track.id}>
                          <Panel size={{ unit: 'px', count: track.height }}>
                            <div className='stem-row h-full'>
                              <MultiTrackTrackHeader trackId={track.id} />
                            </div>
                          </Panel>
                          <Divider influence='neither'>
                            <TrackDivider />
                          </Divider>
                        </Fragment>
                      ))}
                      {/* for minimum bottom padding when content overflows and is scrollable */}
                      <Panel size={smallPanelSize}>
                        <div />
                      </Panel>
                      {/* to fill remaining space when content does not overflow */}
                      <Panel size={oneFRPanelSize}>
                        <div />
                      </Panel>
                    </VerticalPanelGroup>
                  </StemsList>
                </div>
              </Panel>
              <Divider influence='neither'>
                <TrackDivider />
              </Divider>
              <Panel size={smallPanelSize}>
                <MultiTrackHeaderListButton
                  onClick={() => setAddTracksTask('add')}
                >
                  + Add Tracks
                </MultiTrackHeaderListButton>
              </Panel>
            </VerticalPanelGroup>
          </MultiTrackTrackHeaders>
        </MultiTrackTrackSection>
        {stemsContext.stemTracks.length === 0 && (
          <WalkthroughWrapper style={{ top: 180 + fullSongPanelSize.count }}>
            <WalkthroughCard onClick={() => setAddTracksTask('add')}>
              <FourStarIcon className='h-8 w-8' />
              <WalkthroughTitle>Add a New Track</WalkthroughTitle>
              <WalkthroughDescription>
                Generate new instruments or vocals based on the original song.
              </WalkthroughDescription>
            </WalkthroughCard>
            <WalkthroughCard onClick={() => setAddTracksTask('extract')}>
              <ShareArrowIcon className='h-8 w-8' />
              <WalkthroughTitle>Extract Stems</WalkthroughTitle>
              <WalkthroughDescription>
                Extract instrumental or vocal performances from the original
                song.
              </WalkthroughDescription>
            </WalkthroughCard>
            <WalkthroughCard onClick={() => setAddTracksTask('remove')}>
              <BiSolidEraser className='h-8 w-8' />
              <WalkthroughTitle>Remove Instruments</WalkthroughTitle>
              <WalkthroughDescription>
                Create new versions of the original song without specific
                instruments.
              </WalkthroughDescription>
            </WalkthroughCard>
          </WalkthroughWrapper>
        )}
        <CanvasRendererContainer fadeOut={fadeOut} ref={receiveWrapperRef}>
          <ScrollCatcher
            onScroll={receiveScroll}
            onZoom={handleZoom}
            scrollXRef={scrollXRef}
            pxPerBeatRef={pxPerBeatRef}
            invertZoom={false} // TODO
            zoomOnResize
          >
            <CanvasRenderer
              getCanvasRegions={getCanvasRegions}
              frameCountRef={frameCountRef}
            />
          </ScrollCatcher>
        </CanvasRendererContainer>
        <SelectedTrackSidebar />
      </MultiTrackTimelineWrapper>
    </TimelineContainer>
  );
});