import { format } from 'date-fns';

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

import { ExactSecondsRange, InfillPayload } from '../types';
import checkForDumbMistakes from './checkForDumbMistakes';
import { resolveLyrics } from './resolveLyrics';
import resolveTargetClipAndRelativeTimings from './resolveTargetClipAndRelativeTimings';
import resolveTimeRange from './resolveTimeRange';
import snapExactSecondsRange from './snapExactSecondsRange';
import validateContextWindow from './validateContextWindow';

const validateFixedInfillContextWindow = (contextWindow: ExactSecondsRange) => {
  const currentDuration = contextWindow.endSeconds - contextWindow.startSeconds;
  const requiredDuration = 30;
  if (Math.abs(currentDuration - requiredDuration) > 0.02) {
    throw new Error(
      `Fixed infill context window must be exactly ${requiredDuration} seconds long, but is ${currentDuration} seconds long`
    );
  }

  return true;
};

export default async function prepareInfillRequest(
  clipsStore: ClipsStore,
  payload: InfillPayload
): Promise<GenerateParams> {
  const { timing, selection, contextWindow } = payload;
  const projectRelativeSelection = resolveTimeRange(timing, selection);
  const projectRelativeContext = resolveTimeRange(timing, contextWindow);

  validateContextWindow(
    projectRelativeSelection.beats,
    projectRelativeContext.beats
  );

  if (payload.mode === 'fixed') {
    validateFixedInfillContextWindow(projectRelativeContext.seconds);
  }

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

  const {
    clip: targetClip,
    clipId: targetClipId,
    clipRelativeContextRangeSeconds,
    clipRelativeSelectionRangeSeconds,
    alignmentSpec,
  } = await resolveTargetClipAndRelativeTimings(
    clipsStore,
    {
      ...payload,
      projectRelativeContext: projectRelativeContext,
      projectRelativeSelection: projectRelativeSelection,
      originalLyrics: currentContextWindowLyrics,
    },
    payload.mode === 'fixed'
  );

  const snappedSelectionRangeSeconds = snapExactSecondsRange(
    clipRelativeSelectionRangeSeconds
  );

  if (payload.mode === 'fixed') {
    if (targetClip.metadata.duration && targetClip.metadata.duration < 30) {
      throw new Error(
        'Cannot use fixed infill on a clip with duration less than 30 seconds. Try classic mode instead.'
      );
    }
    // Removed the context window adjustment logic that constrained it to clip bounds
    // The context window can now extend beyond song bounds as needed
  }

  const snappedContextRangeSeconds = snapExactSecondsRange(
    clipRelativeContextRangeSeconds,
    payload.mode === 'fixed' ? 30 : undefined
  );

  if (payload.mode === 'fixed') {
    validateFixedInfillContextWindow(snappedContextRangeSeconds);
  }

  const result: GenerateParams = {
    studioProjectId: payload.studioProjectId,
    projectId: payload.projectId || undefined,
    prompt: {
      title: `Replace ${encodeTimeFormat(Math.max(0, projectRelativeSelection.seconds.startSeconds))}-${encodeTimeFormat(Math.max(0, projectRelativeSelection.seconds.endSeconds))} (${format(Date.now(), 'hh:mm a MMM dd')})`,
      type: PromptType.Custom,
      tags: payload.styles,
      negativeTags: payload.excludeStyles,
      lyrics: contextWindowLyrics,
    },
    references: [
      {
        type:
          payload.mode === 'fixed'
            ? ReferenceType.FixedInfill
            : ReferenceType.Infill,
        clipId: targetClipId,
        durationSeconds:
          snappedSelectionRangeSeconds.endSeconds -
          snappedSelectionRangeSeconds.startSeconds,
        startSeconds: snappedSelectionRangeSeconds.startSeconds,
        endSeconds: snappedSelectionRangeSeconds.endSeconds,
        contextStartSeconds: snappedContextRangeSeconds.startSeconds,
        contextEndSeconds: snappedContextRangeSeconds.endSeconds,
        infillLyrics,
        contextLyrics: contextWindowLyrics,
        lyricsUpdated: payload.lyricsUpdated,
      },
    ],
    modelTier: payload.modelTier,
    modelOverride: payload.modelOverride || undefined,
    alignmentOverrides: alignmentSpec,
  };

  if (result.references![0]!.type === ReferenceType.Infill) {
    result.references![0].generateFullContext = true;
  }
  checkForDumbMistakes(payload, result);

  return result;
}