import { GenerateParams, PromptType, ReferenceType } from '@/hooks/useGenerate';

import { getPlaintextLyrics } from '../edit2025/lyrics/getPlaintextLyrics';
import splitLyrics from '../edit2025/lyrics/splitLyrics';
import { AlignedLyric, EditTiming } from '../edit2025/types';
import { StudioContextType } from './StudioContext';
import { getClipContentSeconds } from './getClipContentSeconds';
import { getDerivedTiming, getFocusedStudioClip } from './selectors';
import { StudioClip } from './types';

const replaceEditedLyrics = (
  alignedLyrics: AlignedLyric[],
  timing: EditTiming,
  clip: StudioClip,
  newLyrics: string
) => {
  const splitStart = getClipContentSeconds(clip, timing, clip.startBeats)[0];
  const splitEnd = getClipContentSeconds(clip, timing, clip.endBeats)[1];

  const splitResult = splitLyrics(alignedLyrics, splitStart, splitEnd);

  const result = getPlaintextLyrics([
    ...splitResult[0],
    {
      text: newLyrics,
    },
    ...splitResult[2],
  ]);

  return result;
};

export default async function getUpdatedRegenerateParams(
  regenerateParams: GenerateParams,
  studioContext: StudioContextType,
  alignedLyricsByClipId: Record<string, AlignedLyric[]>
): Promise<GenerateParams> {
  const result = { ...regenerateParams };

  if (!studioContext.studioProjectId) {
    throw new Error('No studio project id');
  }

  result.modelTier = studioContext.modelTier;
  result.projectId = studioContext.projectId;
  result.studioProjectId = studioContext.studioProjectId;

  const focusedStudioClip = getFocusedStudioClip(studioContext.state);
  const timing = getDerivedTiming(studioContext.state);
  const currentClipId = focusedStudioClip?.clipId;
  const currentClipLyrics = currentClipId
    ? alignedLyricsByClipId[currentClipId] || []
    : [];

  if (result.prompt.type === PromptType.Custom) {
    const { styles, excludeStyles } = studioContext.stylesEditController;

    result.prompt = {
      ...result.prompt,
      tags: styles || result.prompt.tags,
      negativeTags: excludeStyles || result.prompt.negativeTags,
    };
  }

  const relevantReferences = (result.references || []).filter(
    (ref) =>
      ref.type === ReferenceType.Infill ||
      ref.type === ReferenceType.FixedInfill ||
      ref.type === ReferenceType.Extend
  );

  if (relevantReferences.length) {
    const lyricsEditSpec =
      studioContext.lyricsEditController.getLyricsEditSpec();
    const infillLyrics = lyricsEditSpec.editedPlaintextLyrics;

    const contextWindowLyrics = focusedStudioClip
      ? replaceEditedLyrics(
          currentClipLyrics,
          timing,
          focusedStudioClip,
          infillLyrics
        )
      : infillLyrics;

    const lyricsUpdated =
      lyricsEditSpec.originalPlaintextLyrics !==
      lyricsEditSpec.editedPlaintextLyrics;

    relevantReferences.forEach((ref) => {
      ref.contextLyrics = contextWindowLyrics;
      if (ref.type !== ReferenceType.Extend) {
        ref.infillLyrics = infillLyrics;
      }
      ref.lyricsUpdated = lyricsUpdated;
    });

    if (result.prompt.type === PromptType.Custom) {
      result.prompt.lyrics = infillLyrics;
    }
  }

  return result;
}