import { addSecondsToBeats } from '@suno/studiokit/timeMapping';
import { format } from 'date-fns';

import { EditTiming } from '@/components/edit2025/types';
import { GenerateParams, PromptType, ReferenceType } from '@/hooks/useGenerate';
import { ClipsStore } from '@/state/clipStore';
import { encodeTimeFormat } from '@/utils/utils';

import { ExactBeatRange, ExtendPayload, StudioTrack } from '../types';
import checkForDumbMistakes from './checkForDumbMistakes';
import { resolveLyrics } from './resolveLyrics';
import resolveTargetClipAndRelativeTimings from './resolveTargetClipAndRelativeTimings';
import resolveTimeRange from './resolveTimeRange';
import snapExactSecondsRange from './snapExactSecondsRange';

const getExtendContext = (
  tracks: StudioTrack[],
  extendFromBeats: number,
  timing: EditTiming
): ExactBeatRange => {
  return {
    type: 'exact-beat-range',
    endBeats: extendFromBeats,
    startBeats: Math.max(
      addSecondsToBeats(-240, extendFromBeats, timing),
      Math.min(...tracks.flatMap((t) => t.clips.map((c) => c.startBeats)))
    ),
  };
};

const getExtendSelection = (extendFromBeats: number): ExactBeatRange => {
  return {
    type: 'exact-beat-range',
    endBeats: Infinity,
    startBeats: extendFromBeats,
  };
};

export default async function prepareExtendRequest(
  clipsStore: ClipsStore,
  payload: ExtendPayload
): Promise<GenerateParams> {
  const { timing, extendFromBeats, tracks } = payload;
  const autoEvaluatedContext = getExtendContext(
    tracks,
    extendFromBeats,
    timing
  );
  const projectRelativeContext = resolveTimeRange(timing, autoEvaluatedContext);
  const projectRelativeSelection = resolveTimeRange(
    timing,
    getExtendSelection(extendFromBeats)
  );

  const { infillLyrics, contextWindowLyrics, currentContextWindowLyrics } =
    resolveLyrics(projectRelativeContext.beats, payload.lyrics);

  const {
    clipId: targetClipId,
    clipRelativeSelectionRangeSeconds,
    alignmentSpec,
  } = await resolveTargetClipAndRelativeTimings(clipsStore, {
    ...payload,
    projectRelativeContext: projectRelativeContext,
    projectRelativeSelection: projectRelativeSelection,
    originalLyrics: currentContextWindowLyrics,
  });

  const snappedSelectionRangeSeconds = snapExactSecondsRange(
    clipRelativeSelectionRangeSeconds
  );

  const result: GenerateParams = {
    studioProjectId: payload.studioProjectId,
    projectId: payload.projectId || undefined,
    prompt: {
      title: `Extend from ${encodeTimeFormat(Math.max(0, projectRelativeSelection.seconds.startSeconds))} (${format(Date.now(), 'hh:mm a MMM dd')})`,
      type: PromptType.Custom,
      tags: payload.styles,
      negativeTags: payload.excludeStyles,
      lyrics: infillLyrics,
    },
    references: [
      {
        type: ReferenceType.Extend,
        clipId: targetClipId,
        startSeconds: snappedSelectionRangeSeconds.startSeconds,
        contextLyrics: contextWindowLyrics,
        isUpload: false,
        lyricsUpdated: payload.lyricsUpdated,
      },
    ],
    modelTier: payload.modelTier,
    modelOverride: payload.modelOverride || undefined,
    alignmentOverrides: alignmentSpec,
  };

  checkForDumbMistakes(payload, result);

  return result;
}