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

import type { LLMModelType, NodeType, SelectedModel, WorkflowNodeData } from "@/types";
import type { ProviderModel } from "@/lib/providers/types";
import { isPopiProviderMode, POPI_PROVIDER_ID } from "@/lib/providerMode";
import { useModelStore, type PopiModelKind } from "@/store/modelStore";
import { loadNodeDefaults } from "@/store/utils/localStorage";
import {
  POPISERVER_DEFAULT_AUDIO_MODEL_ID,
  POPISERVER_DEFAULT_IMAGE_MODEL_ID,
  POPISERVER_DEFAULT_LLM_MODEL_ID,
  POPISERVER_DEFAULT_VIDEO_MODEL_ID,
} from "@/store/utils/defaultImageModel";

type GenerationNodeType = "nanoBanana" | "generateVideo" | "generateAudio" | "llmGenerate";

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

function getConfiguredPopiModelId(type: GenerationNodeType): string | null {
  const defaults = loadNodeDefaults();
  if (type === "nanoBanana") {
    const model = defaults.generateImage?.selectedModel;
    return model?.provider === POPI_PROVIDER_ID && model.modelId ? model.modelId : null;
  }
  if (type === "generateVideo") {
    const model = defaults.generateVideo?.selectedModel;
    return model?.provider === POPI_PROVIDER_ID && model.modelId ? model.modelId : null;
  }
  if (type === "generateAudio") {
    const model = defaults.generateAudio?.selectedModel;
    return model?.provider === POPI_PROVIDER_ID && model.modelId ? model.modelId : null;
  }
  const llm = defaults.llm;
  return llm?.provider === POPI_PROVIDER_ID && llm.model ? llm.model : null;
}

function getCodeDefaultPopiModelId(type: GenerationNodeType): string {
  if (type === "nanoBanana") return POPISERVER_DEFAULT_IMAGE_MODEL_ID;
  if (type === "generateVideo") return POPISERVER_DEFAULT_VIDEO_MODEL_ID;
  if (type === "generateAudio") return POPISERVER_DEFAULT_AUDIO_MODEL_ID;
  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 providerModelToSelectedModel(model: ProviderModel): SelectedModel {
  return {
    provider: model.provider,
    modelId: model.id,
    displayName: model.name,
    capabilities: model.capabilities,
    metadata: model.metadata,
    pricing: model.pricing
      ? {
          type: model.pricing.type,
          amount: model.pricing.amount,
          variants: model.pricing.variants,
        }
      : undefined,
  };
}

export async function buildDefaultGenerationNodeData(
  type: NodeType
): Promise<Partial<WorkflowNodeData> | undefined> {
  if (!isPopiProviderMode()) return undefined;

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

  const generationType = type as GenerationNodeType;
  await useModelStore.getState().initializePopiModels();

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

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

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

  return { selectedModel } as Partial<WorkflowNodeData>;
}