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 } from "@/store/generationPreferenceStore";
import { toSelectedModel } from "@/utils/selectedModel";

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 { type: _type, model: _model, ...preferenceConfig } = preferences.image;
    const aspectRatio = preferences.image.aspectRatio as AspectRatio;
    const selectedModel = resolvePreferenceModel("image", preferences.image.model);
    return {
      ...preferenceConfig,
      aspectRatio,
      resolution: preferences.image.resolution,
      selectedModel,
      imageCount: preferences.image.imageCount,
      parameters: preferences.image.parameters ?? {},
    } as Partial<WorkflowNodeData>;
  }

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

  if (type === "smartAudio" || type === "generateAudio") {
    const { type: _type, model: _model, ...preferenceConfig } = preferences.audio;
    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 {
      ...preferenceConfig,
      selectedModel,
      audioVoiceId: preferences.audio.voiceId,
      parameters,
    } as Partial<WorkflowNodeData>;
  }

  return undefined;
}

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>;
}