import React, { useCallback, useEffect, useMemo, useRef } from 'react';
import { twMerge } from 'tailwind-merge';

import useChoreography, { ChoreographySequence } from '@/hooks/useChoreography';
import { AlignedLyrics } from '@/state/clipStore';

export type HighlightLineOffsetData = {
  alignedLyrics: AlignedLyrics;
  lineIndex: number;
  currentLine: AlignedLyrics[number];
  previousLine?: AlignedLyrics[number];
  nextLine?: AlignedLyrics[number];
  previousLineTime?: number;
  nextLineTime?: number;
  groupTiming?: Array<{ startS: number; endS: number }>;
};

export type HighlightWordOffsetData = HighlightLineOffsetData & {
  currentWord: AlignedLyrics[number]['words'][number];
  previousWord?: AlignedLyrics[number]['words'][number];
  nextWord?: AlignedLyrics[number]['words'][number];
  previousWordTime?: number;
  nextWordTime?: number;
};

export type HighlightLineOffset = (data: HighlightLineOffsetData) => number;
export type HighlightWordOffset = (data: HighlightWordOffsetData) => number;

export enum LyricsHighlightMode {
  /** Highlighted words last the duration of the syllable */
  Syllable = 'syllable',
  /** Highlighted words last the duration of the word */
  Word = 'word',
  /** Highlighted words last the duration of the line */
  Line = 'line',
  /** Highlighted words last the duration of the group */
  Group = 'group',
}

export type LyricsRendererProps = {
  className?: string;
  lineClassName?: string;
  wordClassName?: string;
  /**
   * Current time relative to the aligned lyrics
   *
   * It is the consumer's responsibility to keep track of time and sync with
   * audio or video playback
   */
  currentTime?: number;
  /**
   * If set, changes in `currentTime` that exceed this threshold will invoke
   * each callback from 0
   */
  resetThreshold?: number;
  /** When set, highlight lines in groups rather than individually */
  groupLines?: boolean;
  /** Adjusts the highlight start time of a group */
  groupOffsetStart?: number;
  /** Adjusts the highlight end time of a group */
  groupOffsetEnd?: number;
  /** Maximum duration to consider when grouping */
  maxGroupDuration?: number;
  /** Maximum number of lines to put in a group */
  maxLinesPerGroup?: number;
  /** Naturally break groups on long pauses between lines */
  maxPauseBetweenLines?: number;
  /** When set, don't create groups that cross section boundaries */
  respectSections?: boolean;
  /** When set, don't create groups that cross section boundaries */
  highlightMode?: LyricsHighlightMode;
  /** Adjusts the highlight start time of a line */
  highlightLineOffsetStart?: number | HighlightLineOffset;
  /** Adjusts the highlight end time of a line */
  highlightLineOffsetEnd?: number | HighlightLineOffset;
  /** Adjusts the highlight start time of a word */
  highlightWordOffsetStart?: number | HighlightWordOffset;
  /** Adjusts the highlight end time of a word */
  highlightWordOffsetEnd?: number | HighlightWordOffset;
  /** Lyrics broken down into lines and words with timing information */
  alignedLyrics: AlignedLyrics;
  ref?: React.Ref<HTMLDivElement | null>;
};

export type Props = Omit<
  React.HTMLAttributes<HTMLDivElement>,
  keyof LyricsRendererProps
> &
  LyricsRendererProps;

/**
 * The lyrics aligner sometimes splits words into syllables
 */
export function combineLyricSyllables(alignedLyrics: AlignedLyrics) {
  return alignedLyrics.map((line) => ({
    ...line,
    words: line.words.reduce(
      (lineWords, word) => {
        if (!lineWords.length || word.text.startsWith(' ')) {
          // Add new word
          lineWords.push({ ...word });
        } else {
          // Merge with previous word
          const previousWord = lineWords[lineWords.length - 1];
          previousWord.text += word.text;
          previousWord.endS = word.endS;
        }
        return lineWords;
      },
      [] as typeof line.words
    ),
  }));
}

/**
 * Generates an array of "grouped" timings for the given aligned lyrics
 *
 * The resulting array has the same number of elements as the input, but lines
 * in the same group will have identical start/end times.
 */
export function getGroupedLineTiming(
  alignedLyrics: AlignedLyrics,
  options: Pick<
    LyricsRendererProps,
    | 'groupOffsetStart'
    | 'groupOffsetEnd'
    | 'maxGroupDuration'
    | 'maxLinesPerGroup'
    | 'maxPauseBetweenLines'
    | 'respectSections'
  >
) {
  const {
    maxGroupDuration = 8,
    maxLinesPerGroup = 4,
    maxPauseBetweenLines = 2,
    respectSections = true,
    groupOffsetStart = -2,
    groupOffsetEnd = 2,
  } = options || {};

  // Create groups of lines
  const groups: Array<{
    startS: number;
    endS: number;
    lines: Array<{ line: AlignedLyrics[number]; index: number }>;
  }> = [];

  let currentGroup: Array<{ line: AlignedLyrics[number]; index: number }> = [];
  let currentGroupStartTime = 0;

  alignedLyrics.forEach((line, i) => {
    // Initialize group if empty
    if (currentGroup.length === 0) {
      currentGroupStartTime = Math.max(0, line.startS);
      currentGroup.push({ line, index: i });
      return;
    }

    // Calculate potential group duration if we add this line
    const potentialDuration = line.endS - currentGroupStartTime;

    // Decide whether to add to current group or start new one
    const shouldStartNewGroup =
      // Group would be too long
      (maxGroupDuration && potentialDuration > maxGroupDuration) ||
      // Too many lines
      (maxLinesPerGroup && currentGroup.length >= maxLinesPerGroup) ||
      // Section boundary (if respecting sections)
      (respectSections && currentGroup[0].line.section !== line.section) ||
      // Large pause between lines suggests natural break
      (maxPauseBetweenLines &&
        line.startS - currentGroup[currentGroup.length - 1].line.endS >
          maxPauseBetweenLines);

    if (shouldStartNewGroup) {
      // Finalize current group
      const groupEndTime = currentGroup[currentGroup.length - 1].line.endS;
      groups.push({
        startS: currentGroupStartTime,
        endS: groupEndTime,
        lines: currentGroup,
      });

      // Start new group
      currentGroup = [{ line, index: i }];
      currentGroupStartTime = Math.max(0, line.startS);
    } else {
      // Add to current group
      currentGroup.push({ line, index: i });
    }
  });

  // Don't forget the last group
  if (currentGroup.length > 0) {
    const groupEndTime = currentGroup[currentGroup.length - 1].line.endS;
    groups.push({
      startS: currentGroupStartTime,
      endS: groupEndTime,
      lines: currentGroup,
    });
  }

  // Return simplified timing objects
  return alignedLyrics.map((originalLine, index) => {
    const groupIndex = groups.findIndex((group) =>
      group.lines.some((line) => line.index === index)
    );

    // If it's not in a group, just return as-is
    if (groupIndex < 0) {
      return {
        startS: originalLine.startS,
        endS: originalLine.endS,
      };
    }

    const group = groups[groupIndex];
    const previousGroup = groupIndex > 0 ? groups[groupIndex - 1] : undefined;
    const nextGroup =
      groupIndex < groups.length - 1 ? groups[groupIndex + 1] : undefined;

    const adjustedStartTime =
      previousGroup == null
        ? group.startS + groupOffsetStart
        : Math.max(
            groupOffsetEnd < 0
              ? previousGroup.endS + groupOffsetEnd
              : previousGroup.endS,
            group.startS + groupOffsetStart
          );
    const adjustedEndTime =
      nextGroup == null
        ? group.endS + groupOffsetEnd
        : Math.min(
            Math.max(
              groupOffsetStart < 0
                ? nextGroup.startS + groupOffsetStart
                : nextGroup.startS,
              group.endS
            ),
            group.endS + groupOffsetEnd
          );

    return {
      startS: Math.max(0, adjustedStartTime),
      endS: Math.max(0, adjustedEndTime),
    };
  });
}

export function getGroupedHighlightStart(data: HighlightLineOffsetData) {
  const { lineIndex, currentLine, groupTiming } = data;
  return groupTiming && lineIndex < groupTiming.length
    ? groupTiming[lineIndex].startS
    : currentLine.startS;
}

export function getGroupHighlightEnd(data: HighlightLineOffsetData) {
  const { lineIndex, currentLine, groupTiming } = data;
  return groupTiming && lineIndex < groupTiming.length
    ? groupTiming[lineIndex].endS
    : currentLine.endS;
}

const LyricsRenderer: React.FC<Props> = (props) => {
  const {
    className,
    lineClassName,
    wordClassName: explicitWordClassName,
    currentTime,
    resetThreshold,
    alignedLyrics,
    groupLines = false,
    groupOffsetStart,
    groupOffsetEnd,
    maxGroupDuration,
    maxLinesPerGroup,
    maxPauseBetweenLines,
    respectSections,
    highlightMode = LyricsHighlightMode.Word,
    highlightLineOffsetStart = 0,
    highlightLineOffsetEnd = 0,
    highlightWordOffsetStart = 0,
    highlightWordOffsetEnd = 0,
    ref: propRef,
    ...restProps
  } = props;

  const stateRef = useRef({ currentTime, resetThreshold });
  const containerRef = useRef<HTMLDivElement>(null);
  const setRef = useCallback(
    (el: HTMLDivElement) => {
      containerRef.current = el;

      if (typeof propRef === 'function') {
        propRef(el);
      } else if (propRef) {
        propRef.current = el;
      }
    },
    [propRef]
  );

  const wordClassName = twMerge(
    'data-highlight:text-foreground-primary',
    'data-highlight:duration-0',
    'transition-all duration-500',
    explicitWordClassName
  );

  const toggleLine = useCallback((i: number, shouldHighlight = true) => {
    if (containerRef.current) {
      const el = containerRef.current.querySelector<HTMLElement>(
        `[data-lyric-line="${i}"]`
      );
      if (el) {
        if (shouldHighlight) {
          el.dataset.highlight = '';
        } else {
          delete el.dataset.highlight;
        }
      }
    }
  }, []);

  const toggleWord = useCallback((i: number, shouldHighlight = true) => {
    if (containerRef.current) {
      const el = containerRef.current.querySelector<HTMLElement>(
        `[data-lyric-word="${i}"]`
      );
      if (el) {
        if (shouldHighlight) {
          el.dataset.highlight = '';
        } else {
          delete el.dataset.highlight;
        }
      }
    }
  }, []);

  const [sequence, lyrics] = useMemo<
    [ChoreographySequence, React.ReactElement[]]
  >(() => {
    let wordCount = 0;
    const choreographySequence: ChoreographySequence = [];
    const lyricsLineElements: React.ReactElement[] = [];

    const cleanedLyrics =
      highlightMode === LyricsHighlightMode.Word
        ? combineLyricSyllables(alignedLyrics)
        : alignedLyrics;

    const groupTiming = getGroupedLineTiming(cleanedLyrics, {
      groupOffsetStart,
      groupOffsetEnd,
      maxGroupDuration,
      maxLinesPerGroup,
      maxPauseBetweenLines,
      respectSections,
    });

    // Flatten lines and words in aligned lyrics
    cleanedLyrics.forEach((line, j) => {
      const lyricsWordElements: React.ReactElement[] = [];
      const previousLine = j > 0 ? cleanedLyrics[j - 1] : undefined;
      const nextLine =
        j < cleanedLyrics.length - 1 ? cleanedLyrics[j + 1] : undefined;

      const lineOffsetData = {
        alignedLyrics: cleanedLyrics,
        lineIndex: j,
        currentLine: line,
        previousLine,
        nextLine,
        previousLineTime: previousLine?.endS,
        nextLineTime: nextLine?.startS,
        groupTiming,
      };

      const lineOffsetStart =
        typeof highlightLineOffsetStart === 'function'
          ? highlightLineOffsetStart(lineOffsetData)
          : highlightLineOffsetStart;
      const lineOffsetEnd =
        typeof highlightLineOffsetEnd === 'function'
          ? highlightLineOffsetEnd(lineOffsetData)
          : highlightLineOffsetEnd;

      const lineStart = groupLines
        ? getGroupedHighlightStart(lineOffsetData)
        : line.startS;
      const lineEnd = groupLines
        ? getGroupHighlightEnd(lineOffsetData)
        : line.endS;

      // Start of the line
      choreographySequence.push({
        time: Math.max(0, lineStart + lineOffsetStart),
        callback: () => {
          toggleLine(j, true);
        },
      });

      // Each word in the line
      line.words.forEach((word, k) => {
        const w = wordCount++;

        const previousWord = k > 0 ? line.words[k - 1] : undefined;
        const nextWord =
          k < line.words.length - 1 ? line.words[k + 1] : undefined;
        const wordOffsetData = {
          ...lineOffsetData,
          wordIndex: w,
          currentWord: word,
          previousWord,
          nextWord,
          previousWordTime: previousWord?.endS,
          nextWordTime: nextWord?.startS,
        };

        const wordOffsetStart =
          typeof highlightWordOffsetStart === 'function'
            ? highlightWordOffsetStart(wordOffsetData)
            : highlightWordOffsetStart;
        const wordOffsetEnd =
          typeof highlightWordOffsetEnd === 'function'
            ? highlightWordOffsetEnd(wordOffsetData)
            : highlightWordOffsetEnd;

        // Element for each word
        lyricsWordElements.push(
          <span
            key={`line-${j}-word-${k}`}
            data-lyric-word={w}
            className={wordClassName}
          >
            {word.text}
          </span>
        );

        // Individual word choreography
        const wordStart = word.startS;
        const wordEnd =
          highlightMode === LyricsHighlightMode.Line
            ? line.endS
            : highlightMode === LyricsHighlightMode.Group
              ? lineEnd
              : highlightMode === LyricsHighlightMode.Word
                ? Math.min(line.endS, nextWord?.startS ?? line.endS)
                : word.endS;

        choreographySequence.push(
          {
            time: Math.max(0, wordStart + wordOffsetStart),
            callback: () => {
              toggleWord(w, true);
            },
          },
          {
            time: Math.max(0, wordEnd + wordOffsetEnd),
            callback: () => {
              toggleWord(w, false);
            },
          }
        );
      });

      // End of the line
      choreographySequence.push({
        time: Math.max(0, lineEnd + lineOffsetEnd),
        callback: () => {
          toggleLine(j, false);
        },
      });

      // Build the elements
      lyricsLineElements.push(
        <p key={`lyric-${j}`} className={lineClassName} data-lyric-line={j}>
          {lyricsWordElements}
        </p>
      );
    });
    return [choreographySequence, lyricsLineElements];
  }, [
    alignedLyrics,
    highlightMode,
    highlightLineOffsetStart,
    highlightLineOffsetEnd,
    highlightWordOffsetStart,
    highlightWordOffsetEnd,
    groupLines,
    groupOffsetStart,
    groupOffsetEnd,
    maxGroupDuration,
    maxLinesPerGroup,
    maxPauseBetweenLines,
    respectSections,
    lineClassName,
    wordClassName,
    toggleLine,
    toggleWord,
  ]);

  const [setTime] = useChoreography({ sequence });

  useEffect(() => {
    stateRef.current.resetThreshold = resetThreshold;
  }, [resetThreshold]);

  useEffect(() => {
    const previousTime = stateRef.current.currentTime;
    const shouldReset =
      // Only reset if we know the time
      !(currentTime == null || previousTime == null) &&
      // Reset if we go backwards
      (currentTime < previousTime ||
        // ...or if we're above the reset threshold
        (stateRef.current.resetThreshold != null &&
          Math.abs(currentTime - previousTime) >
            stateRef.current.resetThreshold));

    if (containerRef.current && shouldReset) {
      // Reset when seeking backwards
      for (const el of containerRef.current.querySelectorAll<HTMLElement>(
        '[data-lyric-line], [data-lyric-word]'
      )) {
        delete el.dataset.highlight;
      }
    }
    stateRef.current.currentTime = currentTime;
    if (currentTime != null) {
      setTime(currentTime, shouldReset);
    }
  }, [setTime, currentTime]);

  return (
    <div className={className} {...restProps} ref={setRef}>
      {lyrics}
    </div>
  );
};

export default LyricsRenderer;
