Genarrative/src/components/image-editor/ImageCanvasGenerationDialogModel.ts

import { formatImageSizeValue } from './ImageCanvasEditorModel';
import type {
  CanvasGenerationDialogState,
  CanvasLayer,
  CanvasViewport,
  CharacterAnimationPanelState,
  GenerateDialogState,
  QuickEditPanelState,
  SpecFormValues,
  SpecGenerationType,
} from './ImageCanvasEditorTypes';
import {
  CHARACTER_ANIMATION_DURATION_OPTIONS,
  CHARACTER_FRAME_DISPLAY_SIZE,
  CHARACTER_FRAME_ORIGINAL_SIZE,
  createCanvasLayerReference,
  DEFAULT_ICON_DESCRIPTIONS,
  DEFAULT_IMAGE_MODEL,
  DEFAULT_SPEC_FORM_VALUES,
  EDITOR_IMAGE_DIMENSION_OPTIONS,
  ICON_DESCRIPTION_LIMIT,
  ICON_FRAME_DISPLAY_SIZE,
  ICON_FRAME_ORIGINAL_SIZE,
  SPEC_FRAME_DISPLAY_SIZE,
  SPEC_FRAME_ORIGINAL_SIZE,
} from './ImageCanvasGenerationModel';

type CanvasSize = { width: number; height: number };

function getViewportWorldCenter({
  canvasSize,
  viewport,
}: {
  canvasSize: CanvasSize;
  viewport: CanvasViewport;
}) {
  const safeScale = viewport.scale > 0 ? viewport.scale : 1;
  return {
    x: (canvasSize.width / 2 - viewport.x) / safeScale,
    y: (canvasSize.height / 2 - viewport.y) / safeScale,
  };
}

function resetFailedGenerationDialog(dialog: GenerateDialogState) {
  return {
    ...dialog,
    status: dialog.status === 'failed' ? 'idle' : dialog.status,
    errorMessage: dialog.status === 'failed' ? undefined : dialog.errorMessage,
  };
}

function resolveImageDimensionDefaults(imageModel: string) {
  const dimensionOptions =
    EDITOR_IMAGE_DIMENSION_OPTIONS[
      imageModel as keyof typeof EDITOR_IMAGE_DIMENSION_OPTIONS
    ] ?? EDITOR_IMAGE_DIMENSION_OPTIONS[DEFAULT_IMAGE_MODEL];
  return {
    aspectRatio: dimensionOptions.aspectRatios[0],
    imageSize:
      dimensionOptions.imageSizes.find((size) => size === '1K') ??
      dimensionOptions.imageSizes[0],
  };
}

export function createGenerateDialogDraft({
  canvasSize,
  viewport,
}: {
  canvasSize: CanvasSize;
  viewport: CanvasViewport;
}): Omit<CanvasGenerationDialogState, 'id'> {
  const placeholderWidth = 420;
  const placeholderHeight = 420;
  const worldCenter = getViewportWorldCenter({ canvasSize, viewport });
  return {
    mode: 'generate',
    prompt: '',
    status: 'idle',
    composerOpen: true,
    placeholder: {
      x: worldCenter.x - placeholderWidth / 2,
      y: worldCenter.y - placeholderHeight / 2,
      width: placeholderWidth,
      height: placeholderHeight,
      originalWidth: 2048,
      originalHeight: 2048,
    },
  };
}

export function createSpecDialogDraft({
  canvasSize,
  viewport,
  specType,
}: {
  canvasSize: CanvasSize;
  viewport: CanvasViewport;
  specType: SpecGenerationType;
}): Omit<CanvasGenerationDialogState, 'id'> {
  const worldCenter = getViewportWorldCenter({ canvasSize, viewport });
  return {
    mode: 'spec',
    prompt: '',
    status: 'idle',
    composerOpen: true,
    specType,
    specValues: { ...DEFAULT_SPEC_FORM_VALUES[specType] },
    placeholder: {
      x: worldCenter.x - SPEC_FRAME_DISPLAY_SIZE.width / 2,
      y: worldCenter.y - SPEC_FRAME_DISPLAY_SIZE.height / 2,
      width: SPEC_FRAME_DISPLAY_SIZE.width,
      height: SPEC_FRAME_DISPLAY_SIZE.height,
      originalWidth: SPEC_FRAME_ORIGINAL_SIZE.width,
      originalHeight: SPEC_FRAME_ORIGINAL_SIZE.height,
    },
  };
}

export function createCharacterGenerationDialogDraft({
  canvasSize,
  viewport,
  imageModel,
}: {
  canvasSize: CanvasSize;
  viewport: CanvasViewport;
  imageModel: string;
}): Omit<CanvasGenerationDialogState, 'id'> {
  const worldCenter = getViewportWorldCenter({ canvasSize, viewport });
  const dimensionDefaults = resolveImageDimensionDefaults(imageModel);
  return {
    mode: 'character',
    prompt: '',
    status: 'idle',
    composerOpen: true,
    characterSpecReference: null,
    characterReferences: [],
    imageModel,
    aspectRatio: dimensionDefaults.aspectRatio,
    imageSize: dimensionDefaults.imageSize,
    placeholder: {
      x: worldCenter.x - CHARACTER_FRAME_DISPLAY_SIZE.width / 2,
      y: worldCenter.y - CHARACTER_FRAME_DISPLAY_SIZE.height / 2,
      width: CHARACTER_FRAME_DISPLAY_SIZE.width,
      height: CHARACTER_FRAME_DISPLAY_SIZE.height,
      originalWidth: CHARACTER_FRAME_ORIGINAL_SIZE.width,
      originalHeight: CHARACTER_FRAME_ORIGINAL_SIZE.height,
    },
  };
}

export function createIconGenerationDialogDraft({
  canvasSize,
  viewport,
  imageModel,
}: {
  canvasSize: CanvasSize;
  viewport: CanvasViewport;
  imageModel: string;
}): Omit<CanvasGenerationDialogState, 'id'> {
  const worldCenter = getViewportWorldCenter({ canvasSize, viewport });
  const dimensionDefaults = resolveImageDimensionDefaults(imageModel);
  return {
    mode: 'icon',
    prompt: '',
    status: 'idle',
    composerOpen: true,
    iconSpecReference: null,
    iconDescriptions: [...DEFAULT_ICON_DESCRIPTIONS],
    imageModel,
    aspectRatio: dimensionDefaults.aspectRatio,
    imageSize: dimensionDefaults.imageSize,
    placeholder: {
      x: worldCenter.x - ICON_FRAME_DISPLAY_SIZE.width / 2,
      y: worldCenter.y - ICON_FRAME_DISPLAY_SIZE.height / 2,
      width: ICON_FRAME_DISPLAY_SIZE.width,
      height: ICON_FRAME_DISPLAY_SIZE.height,
      originalWidth: ICON_FRAME_ORIGINAL_SIZE.width,
      originalHeight: ICON_FRAME_ORIGINAL_SIZE.height,
    },
  };
}

export function createEditDialogDraft(
  sourceLayer: CanvasLayer,
): GenerateDialogState {
  return {
    mode: 'edit',
    prompt: sourceLayer.prompt
      ? `${sourceLayer.prompt}，在保持主体结构的基础上优化画面细节`
      : '',
    status: 'idle',
    composerOpen: true,
    sourceLayerId: sourceLayer.id,
  };
}

export function createQuickEditPanelDraft(
  sourceLayer: CanvasLayer,
): QuickEditPanelState {
  return {
    sourceLayerId: sourceLayer.id,
    prompt: '',
    size: formatImageSizeValue(
      sourceLayer.originalWidth,
      sourceLayer.originalHeight,
    ),
    model: sourceLayer.model?.trim() || DEFAULT_IMAGE_MODEL,
    status: 'idle',
  };
}

export function createCharacterAnimationPanelDraft(
  layer: CanvasLayer,
): CharacterAnimationPanelState | null {
  if (layer.assetKind !== 'character') {
    return null;
  }
  return {
    sourceLayerId: layer.id,
    promptText: '',
    resolution: '480p',
    ratio: 'same',
    frameCount: 32,
    durationSeconds: 4,
    status: 'idle',
  };
}

export function assignCharacterSpecReference(
  dialog: GenerateDialogState | null,
  layer: CanvasLayer,
): GenerateDialogState | null {
  return dialog?.mode === 'character'
    ? {
        ...resetFailedGenerationDialog(dialog),
        characterSpecReference: createCanvasLayerReference(layer),
        composerOpen: true,
      }
    : dialog;
}

export function appendCharacterReference(
  dialog: GenerateDialogState | null,
  layer: CanvasLayer,
): GenerateDialogState | null {
  return dialog?.mode === 'character'
    ? {
        ...resetFailedGenerationDialog(dialog),
        characterReferences: [
          ...(dialog.characterReferences ?? []),
          createCanvasLayerReference(layer),
        ],
        composerOpen: true,
      }
    : dialog;
}

export function assignIconSpecReference(
  dialog: GenerateDialogState | null,
  layer: CanvasLayer,
): GenerateDialogState | null {
  if (layer.assetKind !== 'icon-spec') {
    return dialog;
  }
  return dialog?.mode === 'icon'
    ? {
        ...resetFailedGenerationDialog(dialog),
        iconSpecReference: createCanvasLayerReference(layer),
        composerOpen: true,
      }
    : dialog;
}

export function updateSpecFormDialogValue(
  dialog: GenerateDialogState | null,
  key: keyof SpecFormValues,
  value: string,
): GenerateDialogState | null {
  if (dialog?.mode !== 'spec') {
    return dialog;
  }
  const specType = dialog.specType ?? 'custom';
  return {
    ...resetFailedGenerationDialog(dialog),
    specValues: {
      ...DEFAULT_SPEC_FORM_VALUES[specType],
      ...dialog.specValues,
      [key]: value,
    },
  };
}

export function updateIconDescriptionInDialog(
  dialog: GenerateDialogState | null,
  index: number,
  value: string,
): GenerateDialogState | null {
  return dialog?.mode === 'icon'
    ? {
        ...resetFailedGenerationDialog(dialog),
        iconDescriptions: (
          dialog.iconDescriptions ?? DEFAULT_ICON_DESCRIPTIONS
        ).map((description, descriptionIndex) =>
          descriptionIndex === index ? value : description,
        ),
      }
    : dialog;
}

export function appendIconDescriptionToDialog(
  dialog: GenerateDialogState | null,
): GenerateDialogState | null {
  if (dialog?.mode !== 'icon') {
    return dialog;
  }
  const descriptions = dialog.iconDescriptions ?? DEFAULT_ICON_DESCRIPTIONS;
  if (descriptions.length >= ICON_DESCRIPTION_LIMIT) {
    return dialog;
  }
  return {
    ...resetFailedGenerationDialog(dialog),
    iconDescriptions: [...descriptions, ''],
  };
}

export function updateCharacterAnimationDurationPanel(
  panel: CharacterAnimationPanelState | null,
  frameCountValue: string,
): CharacterAnimationPanelState | null {
  const option = CHARACTER_ANIMATION_DURATION_OPTIONS.find(
    (item) => String(item.frameCount) === frameCountValue,
  );
  if (!option || !panel) {
    return panel;
  }
  return {
    ...panel,
    frameCount: option.frameCount,
    durationSeconds: option.durationSeconds,
    status: panel.status === 'failed' ? 'idle' : panel.status,
    errorMessage: panel.status === 'failed' ? undefined : panel.errorMessage,
  };
}

export function hideGeneratedLayerComposerAfterBlur(
  dialog: GenerateDialogState | null,
): GenerateDialogState | null {
  return (dialog?.mode === 'generate' ||
    dialog?.mode === 'spec' ||
    dialog?.mode === 'character' ||
    dialog?.mode === 'icon') &&
    dialog.status !== 'generating'
    ? {
        ...dialog,
        composerOpen: false,
      }
    : dialog;
}

export function closeGenerateComposerDialog(
  dialog: GenerateDialogState | null,
): GenerateDialogState | null {
  return dialog?.mode === 'generate'
    ? {
        ...dialog,
        composerOpen: false,
      }
    : dialog;
}