import type { SelectedModel } from "@/types";

export type ReferenceMediaType = "image" | "video" | "audio";

const VIDEO_INPUT_PARAMETER_NAMES = new Set([
  "video",
  "videos",
  "video_url",
  "video_urls",
  "reference_video",
  "reference_video_url",
  "reference_video_urls",
]);

const AUDIO_INPUT_PARAMETER_NAMES = new Set([
  "voice",
  "voices",
  "audio",
  "audios",
  "audio_url",
  "audio_urls",
  "reference_audio",
  "reference_audio_url",
  "reference_audio_urls",
]);

const IMAGE_INPUT_PARAMETER_NAMES = new Set([
  "image",
  "images",
  "image_url",
  "image_urls",
  "first_frame",
  "first_frame_url",
  "last_frame",
  "last_frame_url",
  "tail_image_url",
  "reference_image",
  "reference_image_url",
  "reference_image_urls",
]);

function hasNamedParameterInput(parameters: Record<string, unknown> | undefined, names: Set<string>): boolean {
  if (!parameters) return false;
  return Object.keys(parameters).some((key) => names.has(key));
}

function metadataBoolean(model: SelectedModel, key: string): boolean | null {
  const value = model.metadata?.[key];
  return typeof value === "boolean" ? value : null;
}

export function modelSupportsReferenceMedia(
  model: SelectedModel,
  mediaType: ReferenceMediaType,
  parameters?: Record<string, unknown>
): boolean {
  if (mediaType === "image") {
    const supportImages = metadataBoolean(model, "isSupportImages");
    if (supportImages !== null) return supportImages;
    if (hasNamedParameterInput(parameters, IMAGE_INPUT_PARAMETER_NAMES)) return true;
    return Boolean(model.capabilities?.some((capability) =>
      capability === "image-to-image" ||
      capability === "image-to-video" ||
      capability === "image-to-3d"
    ));
  }

  if (mediaType === "video") {
    const supportVideos = metadataBoolean(model, "isSupportVideos");
    if (supportVideos !== null) return supportVideos;
    if (hasNamedParameterInput(parameters, VIDEO_INPUT_PARAMETER_NAMES)) return true;
    return Boolean(model.capabilities?.includes("video-to-video"));
  }

  const supportAudios = metadataBoolean(model, "isSupportAudios");
  if (supportAudios !== null) return supportAudios;
  if (hasNamedParameterInput(parameters, AUDIO_INPUT_PARAMETER_NAMES)) return true;
  return Boolean(model.capabilities?.includes("audio-to-video"));
}

export function filterReferenceSubjectList(
  referenceSubjectList: unknown,
  model: SelectedModel,
  parameters?: Record<string, unknown>
): unknown {
  if (!Array.isArray(referenceSubjectList)) return referenceSubjectList;
  return referenceSubjectList.filter((item) => {
    if (!item || typeof item !== "object" || Array.isArray(item)) return false;
    const type = (item as Record<string, unknown>).type;
    if (type !== "image" && type !== "video" && type !== "audio") return true;
    return modelSupportsReferenceMedia(model, type, parameters);
  });
}