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

import { StemPayload } from '../types';
import checkForDumbMistakes from './checkForDumbMistakes';
import resolveTargetClipAndRelativeTimings from './resolveTargetClipAndRelativeTimings';
import resolveTimeRange from './resolveTimeRange';

const DISCARD_CONTEXT = true;

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

  const contextPromise = resolveTargetClipAndRelativeTimings(
    clipsStore,
    {
      ...payload,
      title: undefined,
      tracks: payload.tracks.filter((t) => t.id !== payload.focusedTrackId),
      projectRelativeContext: projectRelativeContext,
      projectRelativeSelection: projectRelativeSelection,
      originalLyrics: '', // TODO: full song lyrics!
    },
    false,
    true
  );

  const selectionPromise = payload.tracks
    .find((t) => t.id === payload.focusedTrackId)
    ?.clips.find(
      (c) =>
        c.startBeats < projectRelativeSelection.beats.endBeats &&
        c.endBeats > projectRelativeSelection.beats.startBeats
    )
    ? resolveTargetClipAndRelativeTimings(
        clipsStore,
        {
          ...payload,
          tracks: payload.tracks.filter((t) => t.id === payload.focusedTrackId),
          projectRelativeContext: projectRelativeSelection, // we only want the model to be concerned with the selection content
          projectRelativeSelection: projectRelativeSelection,
          originalLyrics: '', // TODO: full song lyrics!
        },
        false,
        true
      )
    : Promise.resolve(undefined);

  const [contextResult, selectionResult] = await Promise.all([
    contextPromise,
    selectionPromise,
  ]);

  const contextClipId = contextResult.clipId;
  const selectionClipId = selectionResult?.clipId;

  const result: GenerateParams = {
    studioProjectId: payload.studioProjectId,
    projectId: payload.projectId || undefined,
    prompt: {
      type: PromptType.Custom,
      tags: payload.stemConditionStyles ?? undefined,
      title: payload.title,
      lyrics: payload.stemConditionLyrics ?? undefined,
    },
    references: [],
    modelTier: payload.modelTier,
    modelOverride: payload.modelOverride || undefined,
  };

  result.references!.push({
    type: ReferenceType.StemCondition,
    clipId: contextClipId,
    stemControlTags: payload.stemConditionControlTags,
  });
  if (selectionClipId && payload.selectionConditionMode === 'cover') {
    result.references!.push({
      type: ReferenceType.Cover,
      clipId: selectionClipId,
      startSeconds:
        selectionResult!.clipRelativeSelectionRangeSeconds.startSeconds,
      endSeconds: selectionResult!.clipRelativeSelectionRangeSeconds.endSeconds,
    });
  }
  if (selectionClipId && payload.selectionConditionMode === 'persona') {
    result.references!.push({
      type: ReferenceType.Persona,
      clipId: selectionClipId,
      // intentionally nulling and anying here.
      // we usually want a persona id on these clips and in all other cases it would be a mistake to omit it.
      // in this case, though, it definitely doesn't exist
      personaId: null as any,
    });
  }

  checkForDumbMistakes(payload, result);

  return result;
}