import { getBeatsFromZero } from '@suno/studiokit/timeMapping';
import { useAnimationFrame } from 'framer-motion';
import {
  createContext,
  useCallback,
  useContext,
  useEffect,
  useMemo,
  useRef,
} from 'react';

import EditPlaybackContext from './EditPlaybackContext';
import PreviewClipContext from './PreviewClipContext';
import SelectionContext from './SelectionContext';
import useFrameCountRef from './canvasRenderer/useFrameCountRef';
import EditTimingContext from './queryHooks/EditTimingContext';

export const useEditTimelineContext = () => {
  const frameCountRef = useFrameCountRef();
  const { songEndSeconds } = useContext(PreviewClipContext);

  const timing = useContext(EditTimingContext);

  const playbackContext = useContext(EditPlaybackContext);
  const { selectionStartSeconds, selectionEndSeconds } =
    useContext(SelectionContext);

  const wrapperRef = useRef<HTMLDivElement | null>(null);

  useAnimationFrame(
    useCallback(
      () => playbackContext.playing && frameCountRef.current++,
      [playbackContext.playing, frameCountRef]
    )
  );

  const pxPerBeatRef = useRef(100);
  const scrollXRef = useRef(0);
  const scrollYRef = useRef(0);

  const zoomAboutSelectionOrCenter = useCallback(
    (factor: number) => {
      const timelineWrapper = wrapperRef.current;
      if (!timelineWrapper) {
        return;
      }

      const timelineWidth = timelineWrapper.offsetWidth;
      const timelineScroll = scrollXRef.current;
      const timelineCenter = timelineScroll + timelineWidth / 2;
      const centerBeat = timelineCenter / pxPerBeatRef.current;
      let targetBeat = centerBeat;

      const boundedSelectionStartX = Math.max(
        0,
        Math.min(
          timelineWidth,
          getBeatsFromZero(selectionStartSeconds, timing) *
            pxPerBeatRef.current -
            scrollXRef.current
        )
      );
      const boundedSelectionEndX = Math.max(
        0,
        Math.min(
          timelineWidth,
          getBeatsFromZero(selectionEndSeconds, timing) * pxPerBeatRef.current -
            scrollXRef.current
        )
      );
      const selectionCenterX =
        (boundedSelectionStartX + boundedSelectionEndX) / 2;

      if (selectionCenterX !== 0 && selectionCenterX !== timelineWidth) {
        // selection is on-screen, zoom into it.
        targetBeat =
          (selectionCenterX + scrollXRef.current) / pxPerBeatRef.current;
      }

      pxPerBeatRef.current *= factor;

      const newCenter = targetBeat * pxPerBeatRef.current;
      const pxToScroll = newCenter - timelineCenter;
      scrollXRef.current += pxToScroll;
      frameCountRef.current++;
    },
    [
      frameCountRef,
      scrollXRef,
      selectionStartSeconds,
      selectionEndSeconds,
      timing,
    ]
  );

  const handleZoom = useCallback(
    (xFactor: number, _yFactor: number, xScrollAdjustment: number) => {
      pxPerBeatRef.current *= xFactor;
      frameCountRef.current++;
      scrollXRef.current += xScrollAdjustment;
    },
    [frameCountRef, pxPerBeatRef, scrollXRef]
  );

  const handleScroll = useCallback(
    (x: number, y: number) => {
      scrollXRef.current += x;
      scrollYRef.current += y;
      frameCountRef.current++;
    },
    [frameCountRef, scrollXRef]
  );

  const resetZoom = useCallback(() => {
    const wrapper = wrapperRef.current;
    if (wrapper && timing && songEndSeconds) {
      const songStartBeats = getBeatsFromZero(0, timing);
      const songEndBeats = getBeatsFromZero(songEndSeconds, timing);

      pxPerBeatRef.current =
        (wrapper.clientWidth - 60) / (songEndBeats - songStartBeats);
      scrollXRef.current = songStartBeats * pxPerBeatRef.current - 30;
      frameCountRef.current++;
    }
  }, [songEndSeconds, timing]);

  useEffect(() => {
    resetZoom();
  }, [resetZoom]);

  useEffect(() => {
    frameCountRef.current++;
  }, [selectionStartSeconds]);

  return useMemo(
    () => ({
      wrapperRef,
      frameCountRef,
      pxPerBeatRef,
      scrollXRef,
      scrollYRef,
      zoomAboutSelectionOrCenter,
      handleZoom,
      handleScroll,
      resetZoom,
    }),
    [
      wrapperRef,
      frameCountRef,
      pxPerBeatRef,
      scrollXRef,
      scrollYRef,
      zoomAboutSelectionOrCenter,
      handleZoom,
      handleScroll,
      resetZoom,
    ]
  );
};

const EditTimelineContext = createContext<
  ReturnType<typeof useEditTimelineContext>
>(undefined as never);

export const EditTimelineProvider = ({
  children,
}: {
  children: React.ReactNode;
}) => {
  const editTimelineContext = useEditTimelineContext();
  return (
    <EditTimelineContext.Provider value={editTimelineContext}>
      {children}
    </EditTimelineContext.Provider>
  );
};

export default EditTimelineContext;