Genarrative/src/components/image-editor/ImageCanvasUiAssetExtractionOverlayView.test.tsx

/* @vitest-environment jsdom */

import { fireEvent, render, screen, within } from '@testing-library/react';
import type { ImgHTMLAttributes } from 'react';
import { describe, expect, it, vi } from 'vitest';

import type { CanvasLayer } from './ImageCanvasEditorTypes';
import {
  IMAGE_MODEL_GPT_IMAGE_2,
  IMAGE_MODEL_NANOBANANA2,
} from './ImageCanvasGenerationModel';
import type { UiAssetExtractionState } from './ImageCanvasUiAssetExtractionModel';
import { ImageCanvasUiAssetExtractionOverlayView } from './ImageCanvasUiAssetExtractionOverlayView';

type ResolvedAssetImageMockProps = Omit<
  ImgHTMLAttributes<HTMLImageElement>,
  'src'
> & {
  src?: string | null;
  objectKey?: string | null;
  fallbackSrc?: string | null;
  refreshKey?: string | number | null;
};

vi.mock('../ResolvedAssetImage', () => ({
  ResolvedAssetImage: ({
    src,
    objectKey,
    fallbackSrc,
    refreshKey,
    ...props
  }: ResolvedAssetImageMockProps) => (
    <img
      {...props}
      src={src ?? fallbackSrc ?? ''}
      data-object-key={objectKey ?? ''}
      data-fallback-src={fallbackSrc ?? ''}
      data-refresh-key={refreshKey == null ? '' : String(refreshKey)}
    />
  ),
}));

function createLayer(overrides: Partial<CanvasLayer> = {}): CanvasLayer {
  return {
    id: 'layer-source',
    resourceId: 'resource-source',
    title: '源图',
    src: 'data:image/png;base64,source',
    x: 100,
    y: 120,
    width: 320,
    height: 240,
    originalWidth: 640,
    originalHeight: 480,
    zIndex: 4,
    sourceType: 'uploaded',
    ...overrides,
  };
}

function createState(
  overrides: Partial<UiAssetExtractionState> = {},
): UiAssetExtractionState {
  return {
    sourceLayerId: 'layer-source',
    tool: 'rect',
    model: IMAGE_MODEL_NANOBANANA2,
    marks: [],
    references: [],
    draftMark: null,
    status: 'idle',
    ...overrides,
  };
}

describe('ImageCanvasUiAssetExtractionOverlayView', () => {
  it('renders quick edit tools vertically on the right without extraction actions', () => {
    const handleToolChange = vi.fn();

    render(
      <ImageCanvasUiAssetExtractionOverlayView
        variant="quick-edit"
        sourceLayer={createLayer()}
        viewport={{ x: 10, y: 20, scale: 1 }}
        state={createState()}
        onToolChange={handleToolChange}
        onPointerStart={vi.fn()}
        onPointerMove={vi.fn()}
        onPointerEnd={vi.fn()}
      />,
    );

    const toolbar = screen.getByRole('toolbar', {
      name: '快速编辑框选工具',
    });

    expect(toolbar.className).toContain(
      'image-canvas-editor__ui-extraction-toolbar--right',
    );
    expect(within(toolbar).queryByText('框选素材')).toBeNull();
    expect(within(toolbar).queryByText('0')).toBeNull();
    expect(within(toolbar).queryByRole('button', { name: '提取' })).toBeNull();
    expect(
      within(toolbar).getByRole('button', { name: '矩形框选' }),
    ).toBeTruthy();
    expect(
      within(toolbar).getByRole('button', { name: '椭圆框选' }),
    ).toBeTruthy();
    expect(
      within(toolbar)
        .getByRole('button', { name: '矩形框选' })
        .getAttribute('aria-pressed'),
    ).toBe('true');
    fireEvent.click(within(toolbar).getByRole('button', { name: '矩形框选' }));
    fireEvent.click(
      within(toolbar).getByRole('button', { name: '画笔自由框选' }),
    );

    expect(handleToolChange).toHaveBeenNthCalledWith(1, null);
    expect(handleToolChange).toHaveBeenNthCalledWith(2, 'brush');
  });

  it('renders UI extraction tools on the right and a bottom composer with previews and selected model pricing', () => {
    const marks = Array.from({ length: 7 }, (_, index) => ({
      id: `mark-${index + 1}`,
      tool: 'rect' as const,
      x: index * 8,
      y: index * 6,
      width: 48,
      height: 36,
    }));

    const onRequestUpload = vi.fn();
    const onRemoveReference = vi.fn();

    render(
      <ImageCanvasUiAssetExtractionOverlayView
        sourceLayer={createLayer({ src: 'data:image/png;base64,ui' })}
        viewport={{ x: 10, y: 20, scale: 1 }}
        state={createState({
          marks,
          references: [
            {
              id: 'ui-ref-1',
              label: '参考UI.png',
              src: 'data:image/png;base64,ui-ref',
            },
          ],
        })}
        onToolChange={vi.fn()}
        onPointerStart={vi.fn()}
        onPointerMove={vi.fn()}
        onPointerEnd={vi.fn()}
        onRequestUpload={onRequestUpload}
        onRemoveReference={onRemoveReference}
        onSubmit={vi.fn()}
      />,
    );

    expect(
      screen.getByRole('toolbar', { name: 'UI素材框选工具' }).className,
    ).toContain('image-canvas-editor__ui-extraction-toolbar--right');
    const panel = screen.getByRole('dialog', { name: 'UI素材提取' });
    expect(panel.className).toContain(
      'image-canvas-editor__ui-extraction-panel',
    );
    expect(
      within(panel).getByText('使用框选工具框选你希望从画面中提取的素材'),
    ).toBeTruthy();
    expect(within(panel).getAllByLabelText(/框选区域预览/u)).toHaveLength(7);
    expect(within(panel).getByLabelText('计划规格').textContent).toBe('1:1·2K');
    expect(
      within(panel).getByRole('button', { name: '提取素材模型 nanobanana2' }),
    ).toBeTruthy();
    expect(
      within(panel).queryByRole('button', { name: /提取素材背景色/u }),
    ).toBeNull();
    expect(
      within(panel).queryByRole('button', { name: /提取素材抠图模型/u }),
    ).toBeNull();
    expect(within(panel).getByLabelText('参考UI.png')).toBeTruthy();
    fireEvent.click(
      within(panel).getByRole('button', { name: '上传UI素材参考图' }),
    );
    fireEvent.click(
      within(panel).getByRole('button', { name: '删除参考UI.png' }),
    );
    expect(onRequestUpload).toHaveBeenCalledTimes(1);
    expect(onRemoveReference).toHaveBeenCalledWith('ui-ref-1');
    expect(
      within(panel).queryByRole('button', { name: '取消框选' }),
    ).toBeNull();
    expect(
      within(panel).getByRole('button', { name: '提取' }).textContent,
    ).toContain('24泥点');
  });

  it('changes the UI extraction model from the model menu', () => {
    const onModelChange = vi.fn();

    render(
      <ImageCanvasUiAssetExtractionOverlayView
        sourceLayer={createLayer({ src: 'data:image/png;base64,ui' })}
        viewport={{ x: 10, y: 20, scale: 1 }}
        state={createState({
          marks: [
            {
              id: 'mark-1',
              tool: 'rect',
              x: 10,
              y: 20,
              width: 80,
              height: 60,
            },
          ],
        })}
        onToolChange={vi.fn()}
        onPointerStart={vi.fn()}
        onPointerMove={vi.fn()}
        onPointerEnd={vi.fn()}
        onModelChange={onModelChange}
        onSubmit={vi.fn()}
      />,
    );

    fireEvent.click(
      screen.getByRole('button', { name: '提取素材模型 nanobanana2' }),
    );
    fireEvent.click(
      within(screen.getByRole('menu', { name: '提取素材模型选项' })).getByRole(
        'button',
        { name: 'gpt-image-2' },
      ),
    );

    expect(onModelChange).toHaveBeenCalledWith(IMAGE_MODEL_GPT_IMAGE_2);
  });

  it('disables extraction submission while reference uploads are pending', () => {
    render(
      <ImageCanvasUiAssetExtractionOverlayView
        sourceLayer={createLayer({ src: 'data:image/png;base64,ui' })}
        viewport={{ x: 10, y: 20, scale: 1 }}
        state={createState({
          marks: [
            {
              id: 'mark-1',
              tool: 'rect',
              x: 10,
              y: 20,
              width: 80,
              height: 60,
            },
          ],
        })}
        onToolChange={vi.fn()}
        onPointerStart={vi.fn()}
        onPointerMove={vi.fn()}
        onPointerEnd={vi.fn()}
        onSubmit={vi.fn()}
        hasPendingImageReferenceUploads
      />,
    );

    expect(
      (screen.getByRole('button', { name: '提取' }) as HTMLButtonElement)
        .disabled,
    ).toBe(true);
  });

  it('keeps the current UI extraction model and references when the target model is too small', () => {
    const alertMock = vi.spyOn(window, 'alert').mockImplementation(() => {});
    const onModelChange = vi.fn();

    render(
      <ImageCanvasUiAssetExtractionOverlayView
        sourceLayer={createLayer()}
        viewport={{ x: 10, y: 20, scale: 1 }}
        state={createState({
          references: Array.from({ length: 5 }, (_, index) => ({
            id: `reference-${index + 1}`,
            label: `参考图${index + 1}`,
            src: `/reference-${index + 1}.png`,
          })),
        })}
        onToolChange={vi.fn()}
        onPointerStart={vi.fn()}
        onPointerMove={vi.fn()}
        onPointerEnd={vi.fn()}
        onModelChange={onModelChange}
        onSubmit={vi.fn()}
      />,
    );

    fireEvent.click(
      screen.getByRole('button', { name: '提取素材模型 nanobanana2' }),
    );
    fireEvent.click(
      within(screen.getByRole('menu', { name: '提取素材模型选项' })).getByRole(
        'button',
        { name: 'gpt-image-2' },
      ),
    );

    expect(alertMock).toHaveBeenCalledWith(
      '当前已有 5 张参考图，gpt-image-2 最多允许 4 张，请先删除多余参考图后再切换',
    );
    expect(onModelChange).not.toHaveBeenCalled();
    expect(screen.getByRole('menu', { name: '提取素材模型选项' })).toBeTruthy();
  });

  it('crops mark previews from the same source coordinates as the red selection', () => {
    render(
      <ImageCanvasUiAssetExtractionOverlayView
        sourceLayer={createLayer({ src: 'data:image/png;base64,ui' })}
        viewport={{ x: 10, y: 20, scale: 1 }}
        state={createState({
          marks: [
            {
              id: 'mark-1',
              tool: 'rect',
              x: 160,
              y: 120,
              width: 80,
              height: 60,
            },
          ],
        })}
        onToolChange={vi.fn()}
        onPointerStart={vi.fn()}
        onPointerMove={vi.fn()}
        onPointerEnd={vi.fn()}
        onSubmit={vi.fn()}
      />,
    );

    const preview = screen.getByLabelText('框选区域预览 1');
    const previewImage = within(preview).getByRole('img', {
      name: '源图框选预览 1',
    }) as HTMLImageElement;

    expect(previewImage.getAttribute('src')).toBe('data:image/png;base64,ui');
    expect(previewImage.style.width).toBe('800%');
    expect(previewImage.style.height).toBe('800%');
    expect(previewImage.style.left).toBe('-200%');
    expect(previewImage.style.top).toBe('-200%');
  });

  it('reuses the source layer signed read-url cache for private crop previews', () => {
    render(
      <ImageCanvasUiAssetExtractionOverlayView
        sourceLayer={createLayer({
          src: '/generated-editor-assets/project-1/source.png',
          objectKey: 'generated-editor-assets/project-1/source.png',
          resourceId: 'resource-private',
          taskId: 'task-private',
        })}
        viewport={{ x: 10, y: 20, scale: 1 }}
        state={createState({
          marks: [
            {
              id: 'mark-1',
              tool: 'rect',
              x: 160,
              y: 120,
              width: 80,
              height: 60,
            },
          ],
        })}
        onToolChange={vi.fn()}
        onPointerStart={vi.fn()}
        onPointerMove={vi.fn()}
        onPointerEnd={vi.fn()}
        onSubmit={vi.fn()}
      />,
    );

    const preview = screen.getByLabelText('框选区域预览 1');
    const previewImage = within(preview).getByRole('img', {
      name: '源图框选预览 1',
    });

    expect(previewImage.getAttribute('data-object-key')).toBe(
      'generated-editor-assets/project-1/source.png',
    );
    expect(previewImage.getAttribute('data-refresh-key')).toBe('task-private');
    expect(previewImage.getAttribute('data-fallback-src')).toBe('');
  });

  it('shows numbered labels for quick edit red selections', () => {
    const { container } = render(
      <ImageCanvasUiAssetExtractionOverlayView
        variant="quick-edit"
        showIndexLabels
        sourceLayer={createLayer()}
        viewport={{ x: 0, y: 0, scale: 1 }}
        state={createState({
          marks: [
            {
              id: 'mark-1',
              tool: 'rect',
              x: 20,
              y: 30,
              width: 120,
              height: 90,
            },
            {
              id: 'mark-2',
              tool: 'ellipse',
              x: 220,
              y: 120,
              width: 80,
              height: 60,
            },
          ],
        })}
        onToolChange={vi.fn()}
        onPointerStart={vi.fn()}
        onPointerMove={vi.fn()}
        onPointerEnd={vi.fn()}
      />,
    );

    const labels = Array.from(
      container.querySelectorAll(
        '.image-canvas-editor__ui-extraction-mark-label',
      ),
    ).map((node) => node.textContent);
    const labelCircles = Array.from(
      container.querySelectorAll(
        '.image-canvas-editor__ui-extraction-mark-label circle',
      ),
    ).map((node) => node.getAttribute('r'));

    expect(labels).toEqual(['1', '2']);
    expect(labelCircles).toEqual(['12', '12']);
  });
});