popiart-node-canvas/src/utils/defaultGenerationNodeModel.ts

import type { AspectRatio, LLMModelType, NodeType, SelectedModel, WorkflowNodeData } from "@/types";
import type { ProviderModel } from "@/lib/providers/types";
import { POPI_PROVIDER_ID } from "@/lib/providerMode";
import { useModelStore, type PopiModelKind } from "@/store/modelStore";
import { loadNodeDefaults } from "@/store/utils/localStorage";
import {
  findPopiserverModelByKey,
  POPISERVER_DEFAULT_AUDIO_MODEL_ID,
  POPISERVER_DEFAULT_IMAGE_MODEL_ID,
  POPISERVER_DEFAULT_LLM_MODEL_ID,
  POPISERVER_DEFAULT_VIDEO_MODEL_ID,
} from "@/store/utils/defaultImageModel";
import { useGenerationPreferenceStore, type GenerationPreferenceNodeSize } from "@/store/generationPreferenceStore";
import { toSelectedModel } from "@/utils/selectedModel";
import { filterTaskRoutingParameters } from "@/utils/generationConfig";
import { getGenerationNodeSizeForAspectRatio } from "@/utils/nodeDimensions";

type GenerationNodeType = "smartImage" | "smartVideo" | "smartAudio" | "smartText" | "nanoBanana" | "generateVideo" | "generateAudio" | "llmGenerate";

function getKindForNodeType(type: NodeType): PopiModelKind | null {
  if (type === "smartImage") return "image";
  if (type === "nanoBanana") return "image";
  if (type === "smartVideo") return "video";
  if (type === "generateVideo") return "video";
  if (type === "smartAudio") return "audio";
  if (type === "generateAudio") return "audio";
  if (type === "llmGenerate" || type === "smartText") return "llm";
  return null;
}

function getConfiguredPopiModelId(): string | null {
  const defaults = loadNodeDefaults();
  const llm = defaults.llm;
  return llm?.provider === POPI_PROVIDER_ID && llm.model ? llm.model : null;
}

function getCodeDefaultPopiModelId(): string {
  return POPISERVER_DEFAULT_LLM_MODEL_ID;
}

function matchesConfiguredModel(model: ProviderModel, configuredModelId: string): boolean {
  return (
    model.id === configuredModelId ||
    String(model.metadata?.aiModelId ?? "") === configuredModelId ||
    String(model.metadata?.aiModelCode ?? "") === configuredModelId ||
    String(model.metadata?.aiModelCodeAlias ?? "") === configuredModelId
  );
}

function getDefaultModelIdForPreference(type: "image" | "video" | "audio"): string {
  if (type === "image") return POPISERVER_DEFAULT_IMAGE_MODEL_ID;
  if (type === "video") return POPISERVER_DEFAULT_VIDEO_MODEL_ID;
  return POPISERVER_DEFAULT_AUDIO_MODEL_ID;
}

function resolvePreferenceModel(type: "image" | "video" | "audio", selectedModel: SelectedModel): SelectedModel {
  const models = useModelStore.getState().byKind[type];
  const model =
    findPopiserverModelByKey(models, selectedModel.modelId) ??
    findPopiserverModelByKey(models, selectedModel.displayName) ??
    findPopiserverModelByKey(models, String(selectedModel.aiModelId ?? "")) ??
    findPopiserverModelByKey(models, selectedModel.aiModelCode) ??
    findPopiserverModelByKey(models, selectedModel.aiModelCodeAlias) ??
    findPopiserverModelByKey(models, getDefaultModelIdForPreference(type));

  if (!model) return selectedModel;
  const resolved = toSelectedModel(model);
  const preferences = useGenerationPreferenceStore.getState();
  if (type === "image") preferences.updateImagePreference({ model: resolved });
  if (type === "video") preferences.updateVideoPreference({ model: resolved });
  if (type === "audio") preferences.updateAudioPreference({ model: resolved });
  return resolved;
}

async function buildGenerationPreferenceNodeData(type: NodeType): Promise<Partial<WorkflowNodeData> | undefined> {
  const preferences = useGenerationPreferenceStore.getState();

  if (type === "smartImage" || type === "nanoBanana") {
    const aspectRatio = preferences.image.aspectRatio as AspectRatio;
    const selectedModel = resolvePreferenceModel("image", preferences.image.model);
    return {
      aspectRatio,
      resolution: preferences.image.resolution,
      selectedModel,
      batchSize: preferences.image.batchSize,
      extraTaskParams: preferences.image.extraTaskParams ?? {},
      subType: preferences.image.subType,
      parameters: filterTaskRoutingParameters(preferences.image.parameters as Record<string, unknown> | undefined),
    } as Partial<WorkflowNodeData>;
  }

  if (type === "smartVideo" || type === "generateVideo") {
    const selectedModel = resolvePreferenceModel("video", preferences.video.model);
    return {
      selectedModel,
      batchSize: preferences.video.batchSize,
      extraTaskParams: preferences.video.extraTaskParams ?? {},
      subType: preferences.video.subType ?? selectedModel.subType,
      parameters: filterTaskRoutingParameters(preferences.video.parameters),
    } as Partial<WorkflowNodeData>;
  }

  if (type === "smartAudio" || type === "generateAudio") {
    const parameters = preferences.audio.parameters && typeof preferences.audio.parameters === "object"
      ? preferences.audio.parameters as Record<string, unknown>
      : {};
    const selectedModel = resolvePreferenceModel("audio", preferences.audio.model);
    return {
      selectedModel,
      batchSize: preferences.audio.batchSize,
      extraTaskParams: preferences.audio.extraTaskParams ?? {},
      subType: preferences.audio.subType,
      audioVoiceId: preferences.audio.voiceId,
      parameters: filterTaskRoutingParameters(parameters),
    } as Partial<WorkflowNodeData>;
  }

  return undefined;
}

export function buildDefaultGenerationNodeSize(type: NodeType): GenerationPreferenceNodeSize | undefined {
  const preferences = useGenerationPreferenceStore.getState();
  if (type === "smartImage" || type === "nanoBanana") return preferences.image.nodeSize;
  if (type === "smartVideo" || type === "generateVideo") return preferences.video.nodeSize;
  if (type === "smartAudio" || type === "generateAudio") return preferences.audio.nodeSize;
  return undefined;
}

function readPreferredAspectRatio(data: Partial<WorkflowNodeData> | undefined): unknown {
  if (!data) return null;
  const directAspectRatio = (data as { aspectRatio?: unknown }).aspectRatio;
  if (typeof directAspectRatio === "string" && directAspectRatio.trim()) return directAspectRatio;
  const parameters = (data as { parameters?: unknown }).parameters;
  if (!parameters || typeof parameters !== "object") return null;
  const record = parameters as Record<string, unknown>;
  return (
    record.ratio ??
    record.aspectRatio ??
    record.aspect_ratio ??
    record.videoRatio ??
    record.video_ratio ??
    record.imageRatio ??
    record.image_ratio ??
    null
  );
}

export async function buildDefaultGenerationNodeCreation(type: NodeType): Promise<{
  data: Partial<WorkflowNodeData> | undefined;
  size: GenerationPreferenceNodeSize | undefined;
}> {
  const data = await buildDefaultGenerationNodeData(type);
  const aspectRatio = readPreferredAspectRatio(data);
  return {
    data,
    size: getGenerationNodeSizeForAspectRatio(aspectRatio) ? undefined : buildDefaultGenerationNodeSize(type),
  };
}

export async function buildDefaultGenerationNodeData(
  type: NodeType
): Promise<Partial<WorkflowNodeData> | undefined> {
  const generationPreferenceData = await buildGenerationPreferenceNodeData(type);
  if (generationPreferenceData) return generationPreferenceData;

  const kind = getKindForNodeType(type);
  if (!kind || kind === "3d") return undefined;

  const generationType = type as GenerationNodeType;
  const models = useModelStore.getState().byKind[kind];
  if (models.length === 0) return undefined;

  const preferredModelId = getConfiguredPopiModelId() ?? getCodeDefaultPopiModelId();
  const model = models.find((item) => matchesConfiguredModel(item, preferredModelId)) ?? models[0];
  const selectedModel = toSelectedModel(model);

  if (generationType === "llmGenerate" || generationType === "smartText") {
    return {
      provider: POPI_PROVIDER_ID,
      model: selectedModel.modelId as LLMModelType,
      selectedModel,
    } as Partial<WorkflowNodeData>;
  }

  return { selectedModel } as Partial<WorkflowNodeData>;
}