popiart-node-canvas/src/store/generationPreferenceStore.ts

import { create } from "zustand";
import type { GenerationBatchSize, SelectedModel } from "@/types";
import {
  createPopiserverDefaultAudioModel,
  createPopiserverDefaultImageModel,
  createPopiserverDefaultVideoModel,
} from "@/store/utils/defaultImageModel";

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

export interface GenerationPreferenceNodeSize {
  width: number;
  height: number;
}

export interface ImageGenerationPreference {
  type: "image";
  model: SelectedModel;
  aspectRatio: string;
  resolution: string;
  batchSize: GenerationBatchSize;
  parameters?: Record<string, unknown>;
  extraTaskParams?: Record<string, unknown>;
  subType?: number;
  nodeSize?: GenerationPreferenceNodeSize;
}

export interface VideoGenerationPreference {
  type: "video";
  model: SelectedModel;
  parameters?: Record<string, unknown>;
  extraTaskParams?: Record<string, unknown>;
  subType?: number;
  batchSize: GenerationBatchSize;
  nodeSize?: GenerationPreferenceNodeSize;
}

export interface AudioGenerationPreference {
  type: "audio";
  model: SelectedModel;
  voiceId?: string;
  parameters?: Record<string, unknown>;
  extraTaskParams?: Record<string, unknown>;
  subType?: number;
  batchSize?: GenerationBatchSize;
  nodeSize?: GenerationPreferenceNodeSize;
}

type GenerationPreferencePatch = {
  model?: SelectedModel;
  selectedModel?: SelectedModel;
  aspectRatio?: unknown;
  resolution?: unknown;
  parameters?: unknown;
  extraTaskParams?: unknown;
  subType?: unknown;
  batchSize?: unknown;
  voiceId?: unknown;
  audioVoiceId?: unknown;
  nodeSize?: unknown;
};

interface GenerationPreferenceState {
  image: ImageGenerationPreference;
  video: VideoGenerationPreference;
  audio: AudioGenerationPreference;
  updateImagePreference: (patch: GenerationPreferencePatch) => void;
  updateVideoPreference: (patch: GenerationPreferencePatch) => void;
  updateAudioPreference: (patch: GenerationPreferencePatch) => void;
}

function isRecord(value: unknown): value is Record<string, unknown> {
  return Boolean(value && typeof value === "object" && !Array.isArray(value));
}

function normalizeCommonPreferencePatch(patch: GenerationPreferencePatch) {
  const next: {
    model?: SelectedModel;
    parameters?: Record<string, unknown>;
    extraTaskParams?: Record<string, unknown>;
    subType?: number;
    batchSize?: GenerationBatchSize;
    nodeSize?: GenerationPreferenceNodeSize;
  } = {};
  const model = patch.model ?? patch.selectedModel;
  if (model) next.model = model;
  if (isRecord(patch.parameters)) next.parameters = { ...patch.parameters };
  if (isRecord(patch.extraTaskParams)) next.extraTaskParams = { ...patch.extraTaskParams };
  if (typeof patch.subType === "number") next.subType = patch.subType;
  if (typeof patch.batchSize === "number") next.batchSize = patch.batchSize;
  if (isRecord(patch.nodeSize)) {
    const width = patch.nodeSize.width;
    const height = patch.nodeSize.height;
    if (typeof width === "number" && Number.isFinite(width) && width > 0 &&
      typeof height === "number" && Number.isFinite(height) && height > 0) {
      next.nodeSize = { width, height };
    }
  }
  return next;
}

function normalizeImagePreferencePatch(patch: GenerationPreferencePatch): Partial<Omit<ImageGenerationPreference, "type">> {
  const next: Partial<Omit<ImageGenerationPreference, "type">> = normalizeCommonPreferencePatch(patch);
  if (typeof patch.aspectRatio === "string") next.aspectRatio = patch.aspectRatio;
  if (typeof patch.resolution === "string") next.resolution = patch.resolution;
  return next;
}

function normalizeVideoPreferencePatch(patch: GenerationPreferencePatch): Partial<Omit<VideoGenerationPreference, "type">> {
  return normalizeCommonPreferencePatch(patch);
}

function normalizeAudioPreferencePatch(patch: GenerationPreferencePatch): Partial<Omit<AudioGenerationPreference, "type">> {
  const next: Partial<Omit<AudioGenerationPreference, "type">> = normalizeCommonPreferencePatch(patch);
  if (Object.prototype.hasOwnProperty.call(patch, "voiceId")) {
    next.voiceId = typeof patch.voiceId === "string" ? patch.voiceId : undefined;
  } else if (Object.prototype.hasOwnProperty.call(patch, "audioVoiceId")) {
    next.voiceId = typeof patch.audioVoiceId === "string" ? patch.audioVoiceId : undefined;
  }
  return next;
}

function createInitialImagePreference(): ImageGenerationPreference {
  return {
    type: "image",
    model: createPopiserverDefaultImageModel(),
    aspectRatio: "9:16",
    resolution: "1K",
    batchSize: 1,
  };
}

function createInitialVideoPreference(): VideoGenerationPreference {
  const model = createPopiserverDefaultVideoModel();
  return {
    type: "video",
    model,
    parameters: {},
    extraTaskParams: {},
    subType: model.subType,
    batchSize: 1,
  };
}

function createInitialAudioPreference(): AudioGenerationPreference {
  return {
    type: "audio",
    model: createPopiserverDefaultAudioModel(),
  };
}

export function createInitialGenerationPreferences(): Pick<GenerationPreferenceState, "image" | "video" | "audio"> {
  return {
    image: createInitialImagePreference(),
    video: createInitialVideoPreference(),
    audio: createInitialAudioPreference(),
  };
}

export const useGenerationPreferenceStore = create<GenerationPreferenceState>()(
  (set) => ({
    ...createInitialGenerationPreferences(),
    updateImagePreference: (patch) => set((state) => ({
      image: { ...state.image, ...normalizeImagePreferencePatch(patch), type: "image" },
    })),
    updateVideoPreference: (patch) => set((state) => ({
      video: { ...state.video, ...normalizeVideoPreferencePatch(patch), type: "video" },
    })),
    updateAudioPreference: (patch) => set((state) => ({
      audio: { ...state.audio, ...normalizeAudioPreferencePatch(patch), type: "audio" },
    })),
  })
);