import type { ModelCapability } from "@/lib/providers/types";

export type PopiModelMediaKind = "image" | "video" | "audio" | "llm";

type PopiSubtypeConfig = {
  kind: PopiModelMediaKind;
  capabilities: ModelCapability[];
};

export const POPI_SUBTYPE_CLASSIFICATION: Record<number, PopiSubtypeConfig> = {
  102: { kind: "image", capabilities: ["text-to-image"] },
  103: { kind: "image", capabilities: ["image-to-image"] },
  104: { kind: "image", capabilities: ["image-to-image"] },
  105: { kind: "image", capabilities: ["image-to-image"] },
  106: { kind: "image", capabilities: ["image-to-image"] },
  107: { kind: "image", capabilities: ["image-to-image"] },
  108: { kind: "image", capabilities: ["image-to-image"] },
  109: { kind: "image", capabilities: ["image-to-image"] },
  202: { kind: "video", capabilities: ["image-to-video"] },
  203: { kind: "video", capabilities: ["image-to-video"] },
  204: { kind: "video", capabilities: ["image-to-video"] },
  205: { kind: "video", capabilities: ["image-to-video"] },
  301: { kind: "audio", capabilities: ["text-to-audio"] },
  501: { kind: "llm", capabilities: ["text-to-text"] },
};

export const POPI_SUBTYPES_BY_KIND: Record<PopiModelMediaKind, number[]> = {
  image: [102, 103, 104, 105, 106, 107, 108, 109],
  video: [202, 203, 204, 205],
  audio: [301],
  llm: [501],
};

export const POPI_CAPABILITIES_BY_KIND: Record<PopiModelMediaKind, ModelCapability[]> = {
  image: ["text-to-image", "image-to-image"],
  video: ["text-to-video", "image-to-video", "video-to-video", "audio-to-video"],
  audio: ["text-to-audio"],
  llm: ["text-to-text"],
};

export function getPopiSubTypeFilter(kind: PopiModelMediaKind): string {
  return POPI_SUBTYPES_BY_KIND[kind].join(",");
}

export function getPopiCapabilitiesForSubType(subType: number): ModelCapability[] {
  return POPI_SUBTYPE_CLASSIFICATION[subType]?.capabilities ?? [];
}

export function getPopiKindForSubType(subType: number): PopiModelMediaKind | null {
  return POPI_SUBTYPE_CLASSIFICATION[subType]?.kind ?? null;
}

export function getPopiCapabilitiesForKind(kind: PopiModelMediaKind): ModelCapability[] {
  return POPI_CAPABILITIES_BY_KIND[kind];
}