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

import type { CharacterReferenceImage } from './ImageCanvasEditorTypes';

export type UiAssetExtractionTool = 'rect' | 'ellipse' | 'brush';

export type UiAssetExtractionPoint = {
  x: number;
  y: number;
};

export type UiAssetExtractionBoxMark = {
  id: string;
  tool: 'rect' | 'ellipse';
  x: number;
  y: number;
  width: number;
  height: number;
};

export type UiAssetExtractionBrushMark = {
  id: string;
  tool: 'brush';
  points: UiAssetExtractionPoint[];
};

export type UiAssetExtractionMark =
  | UiAssetExtractionBoxMark
  | UiAssetExtractionBrushMark;

export type UiAssetExtractionState = {
  sourceLayerId: string;
  tool: UiAssetExtractionTool | null;
  model: string;
  marks: UiAssetExtractionMark[];
  references: CharacterReferenceImage[];
  draftMark: UiAssetExtractionMark | null;
  status: 'idle' | 'extracting' | 'failed';
  errorMessage?: string;
};

export const UI_ASSET_EXTRACTION_ASPECT_RATIO = '1:1';
export const UI_ASSET_EXTRACTION_SMALL_IMAGE_SIZE = '1K';
export const UI_ASSET_EXTRACTION_LARGE_IMAGE_SIZE = '2K';
export const UI_ASSET_EXTRACTION_LARGE_SIZE_MARK_THRESHOLD = 6;

export type UiAssetExtractionImageSize =
  | typeof UI_ASSET_EXTRACTION_SMALL_IMAGE_SIZE
  | typeof UI_ASSET_EXTRACTION_LARGE_IMAGE_SIZE;

export function createUiAssetExtractionState(
  sourceLayerId: string,
  options: {
    initialTool?: UiAssetExtractionTool | null;
    model: string;
  },
): UiAssetExtractionState {
  return {
    sourceLayerId,
    tool: options.initialTool === undefined ? 'rect' : options.initialTool,
    model: options.model,
    marks: [],
    references: [],
    draftMark: null,
    status: 'idle',
  };
}

export function resolveUiAssetExtractionGenerationPlan(markCount: number): {
  aspectRatio: typeof UI_ASSET_EXTRACTION_ASPECT_RATIO;
  imageSize: UiAssetExtractionImageSize;
} {
  return {
    aspectRatio: UI_ASSET_EXTRACTION_ASPECT_RATIO,
    imageSize:
      markCount > UI_ASSET_EXTRACTION_LARGE_SIZE_MARK_THRESHOLD
        ? UI_ASSET_EXTRACTION_LARGE_IMAGE_SIZE
        : UI_ASSET_EXTRACTION_SMALL_IMAGE_SIZE,
  };
}

export function createUiAssetExtractionDraftMark({
  tool,
  id,
  point,
}: {
  tool: UiAssetExtractionTool;
  id: string;
  point: UiAssetExtractionPoint;
}): UiAssetExtractionMark {
  if (tool === 'brush') {
    return { id, tool, points: [point] };
  }
  return {
    id,
    tool,
    x: point.x,
    y: point.y,
    width: 0,
    height: 0,
  };
}

export function updateUiAssetExtractionDraftMark(
  draftMark: UiAssetExtractionMark,
  point: UiAssetExtractionPoint,
): UiAssetExtractionMark {
  if (draftMark.tool === 'brush') {
    return {
      ...draftMark,
      points: [...draftMark.points, point],
    };
  }
  return {
    ...draftMark,
    width: point.x - draftMark.x,
    height: point.y - draftMark.y,
  };
}

export function normalizeUiAssetExtractionMark(
  mark: UiAssetExtractionMark,
): UiAssetExtractionMark | null {
  if (mark.tool === 'brush') {
    const points = mark.points.filter(
      (point) => Number.isFinite(point.x) && Number.isFinite(point.y),
    );
    return points.length >= 2 ? { ...mark, points } : null;
  }
  if (
    !Number.isFinite(mark.x) ||
    !Number.isFinite(mark.y) ||
    !Number.isFinite(mark.width) ||
    !Number.isFinite(mark.height)
  ) {
    return null;
  }
  const left = Math.min(mark.x, mark.x + mark.width);
  const top = Math.min(mark.y, mark.y + mark.height);
  const width = Math.abs(mark.width);
  const height = Math.abs(mark.height);
  if (width < 4 || height < 4) {
    return null;
  }
  return {
    ...mark,
    x: left,
    y: top,
    width,
    height,
  };
}

export function buildUiAssetExtractionBrushPath(
  mark: UiAssetExtractionBrushMark,
) {
  const [firstPoint, ...restPoints] = mark.points;
  if (!firstPoint) {
    return '';
  }
  return [
    `M ${firstPoint.x} ${firstPoint.y}`,
    ...restPoints.map((point) => `L ${point.x} ${point.y}`),
  ].join(' ');
}