Genarrative/src/components/image-editor/EditorAgentConversation/useConversationAttachments.ts

import { useCallback, useMemo, useRef, useState } from 'react';

import {
  createEditorAgentAttachmentRef,
  EDITOR_AGENT_MAX_ATTACHMENTS,
  type EditorAgentAttachmentRef,
} from '@/packages/shared/src/contracts';
import type {
  CanvasLayer,
  EditorAsset,
} from '@/src/components/image-editor/ImageCanvasEditorTypes.ts';
import { probeImageFileDimensions } from '@/src/components/image-editor/ImageCanvasFileModel.ts';
import { uploadEditorMediaAssetFile } from '@/src/services/image-editor/editorMediaAssetUploadClient.ts';
import { createEditorProjectResource } from '@/src/services/image-editor/editorProjectClient.ts';

import {
  attachmentKey,
  contextAssetMediaSrc,
  type EditorAgentContextAsset,
} from './common.ts';

export type AttachmentPickerTab = 'canvas' | 'library';

export type EditorAgentAttachmentOption = {
  key: string;
  sourceLabel: string;
  attachment: EditorAgentAttachmentRef;
};

type AttachmentUpdater =
  | EditorAgentAttachmentRef[]
  | ((
      currentAttachments: EditorAgentAttachmentRef[],
    ) => EditorAgentAttachmentRef[]);

function isImageLayer(layer: CanvasLayer) {
  return (
    (layer.mediaType ?? 'image') === 'image' &&
    Boolean(layer.resourceId?.trim()) &&
    layer.src.trim()
  );
}

function isImageAsset(asset: EditorAsset) {
  return (asset.mediaType ?? 'image') === 'image' && asset.src.trim();
}

function createCanvasAttachmentOptions(
  layers: CanvasLayer[],
): EditorAgentAttachmentOption[] {
  return layers.filter(isImageLayer).map((layer) => {
    const referenceId = layer.resourceId || layer.id;
    const attachment = createEditorAgentAttachmentRef({
      source: 'canvas_resource',
      referenceId,
      objectKey: layer.objectKey ?? null,
      imageSrc: layer.src,
      thumbnailSrc: layer.thumbnailSrc ?? null,
      label: layer.title,
      width: layer.width,
      height: layer.height,
    });
    return {
      key: attachmentKey(attachment),
      sourceLabel: '画布',
      attachment,
    };
  });
}

function createLibraryAttachmentOptions(
  assets: EditorAsset[],
): EditorAgentAttachmentOption[] {
  return assets.filter(isImageAsset).map((asset) => {
    const attachment = createEditorAgentAttachmentRef({
      source: 'library_asset',
      referenceId: asset.id,
      objectKey: asset.objectKey ?? null,
      imageSrc: asset.src,
      thumbnailSrc: asset.thumbnailSrc ?? null,
      label: asset.label,
      width: asset.width,
      height: asset.height,
    });
    return {
      key: attachmentKey(attachment),
      sourceLabel: '素材库',
      attachment,
    };
  });
}

function mergeAttachments(
  currentAttachments: EditorAgentAttachmentRef[],
  nextAttachments: EditorAgentAttachmentRef[],
) {
  const merged = [...currentAttachments];
  const existingKeys = new Set(currentAttachments.map(attachmentKey));
  nextAttachments.forEach((attachment) => {
    const key = attachmentKey(attachment);
    if (!existingKeys.has(key)) {
      existingKeys.add(key);
      merged.push(attachment);
    }
  });
  return merged;
}

function extractClipboardImageFiles(clipboardData: DataTransfer | null) {
  if (!clipboardData) {
    return [];
  }
  return Array.from(clipboardData.files ?? []).filter((file) =>
    file.type.startsWith('image/'),
  );
}

async function createPastedAgentImageAttachment(
  projectId: string,
  file: File,
): Promise<EditorAgentAttachmentRef> {
  const [upload, dimensions] = await Promise.all([
    uploadEditorMediaAssetFile(file, 'image', {
      pathSegments: ['editor', 'agent-paste', 'image', `${Date.now()}`],
      entityId: projectId,
      metadata: {
        source: 'agent-input-paste',
      },
    }),
    probeImageFileDimensions(file),
  ]);
  const width = dimensions?.width ?? 1;
  const height = dimensions?.height ?? 1;
  const resource = await createEditorProjectResource(projectId, {
    imageSrc: upload.src,
    objectKey: upload.objectKey,
    assetObjectId: upload.assetObjectId,
    width,
    height,
    sourceType: 'uploaded',
  });
  return createEditorAgentAttachmentRef({
    source: 'canvas_resource',
    referenceId: resource.resourceId,
    objectKey: resource.objectKey ?? upload.objectKey,
    imageSrc: resource.imageSrc,
    thumbnailSrc: null,
    // pasted file.name does not give more information
    label: resource.label ?? '粘贴图片',
    width: resource.width,
    height: resource.height,
  });
}

export function useConversationAttachments({
  projectId,
  layers,
  assets,
}: {
  projectId: string | null;
  layers: CanvasLayer[];
  assets: EditorAsset[];
}) {
  const [attachments, setAttachments] = useState<EditorAgentAttachmentRef[]>(
    [],
  );
  const attachmentsRef = useRef<EditorAgentAttachmentRef[]>([]);
  const [attachmentPickerOpen, setAttachmentPickerOpen] = useState(false);
  const [attachmentPickerTab, setAttachmentPickerTab] =
    useState<AttachmentPickerTab>('canvas');
  const [attachmentError, setAttachmentError] = useState<string | null>(null);
  const [isPastingAttachment, setIsPastingAttachment] = useState(false);
  const isPastingAttachmentRef = useRef(false);
  const [draftAttachmentKeys, setDraftAttachmentKeys] = useState<Set<string>>(
    () => new Set(),
  );

  const canvasAttachmentOptions = useMemo(
    () => createCanvasAttachmentOptions(layers),
    [layers],
  );
  const libraryAttachmentOptions = useMemo(
    () => createLibraryAttachmentOptions(assets),
    [assets],
  );
  const attachmentOptions = useMemo(
    () => [...canvasAttachmentOptions, ...libraryAttachmentOptions],
    [canvasAttachmentOptions, libraryAttachmentOptions],
  );
  const attachmentOptionsByKey = useMemo(() => {
    const optionMap = new Map<string, EditorAgentAttachmentOption>();
    attachmentOptions.forEach((option) => optionMap.set(option.key, option));
    return optionMap;
  }, [attachmentOptions]);

  const updateAttachments = useCallback((updater: AttachmentUpdater) => {
    const nextAttachments =
      typeof updater === 'function' ? updater(attachmentsRef.current) : updater;
    attachmentsRef.current = nextAttachments;
    setAttachments(nextAttachments);
    return nextAttachments;
  }, []);

  const appendAttachments = useCallback(
    (nextAttachments: EditorAgentAttachmentRef[]) => {
      const mergedAttachments = mergeAttachments(
        attachmentsRef.current,
        nextAttachments,
      );
      if (mergedAttachments.length > EDITOR_AGENT_MAX_ATTACHMENTS) {
        setAttachmentError(`最多 ${EDITOR_AGENT_MAX_ATTACHMENTS} 张`);
        return false;
      }
      updateAttachments(mergedAttachments);
      setAttachmentError(null);
      return true;
    },
    [updateAttachments],
  );

  const openAttachmentPicker = useCallback(() => {
    setAttachmentError(null);
    setDraftAttachmentKeys(new Set(attachmentsRef.current.map(attachmentKey)));
    setAttachmentPickerOpen(true);
  }, []);

  const closeAttachmentPicker = useCallback(() => {
    setAttachmentPickerOpen(false);
  }, []);

  const toggleAttachmentKey = useCallback((key: string) => {
    setDraftAttachmentKeys((currentKeys) => {
      const nextKeys = new Set(currentKeys);
      if (nextKeys.has(key)) {
        nextKeys.delete(key);
      } else {
        nextKeys.add(key);
      }
      return nextKeys;
    });
  }, []);

  const applyAttachmentSelection = useCallback(() => {
    const nextAttachments = Array.from(draftAttachmentKeys)
      .map((key) => attachmentOptionsByKey.get(key)?.attachment)
      .filter((attachment): attachment is EditorAgentAttachmentRef =>
        Boolean(attachment),
      );
    if (nextAttachments.length > EDITOR_AGENT_MAX_ATTACHMENTS) {
      setAttachmentError(`最多 ${EDITOR_AGENT_MAX_ATTACHMENTS} 张`);
      return;
    }
    updateAttachments(nextAttachments);
    setAttachmentPickerOpen(false);
  }, [attachmentOptionsByKey, draftAttachmentKeys, updateAttachments]);

  const referenceContextAsset = useCallback(
    (asset: EditorAgentContextAsset) => {
      if (asset.kind === 'attachment') {
        const directAttachment = createEditorAgentAttachmentRef({
          source: asset.source,
          referenceId: asset.referenceId,
          objectKey: asset.objectKey,
          imageSrc: asset.imageSrc,
          thumbnailSrc: asset.thumbnailSrc,
          label: asset.label,
          width: asset.width,
          height: asset.height,
        });
        return directAttachment.referenceId.trim() &&
          directAttachment.imageSrc.trim()
          ? appendAttachments([directAttachment])
          : false;
      }
      const objectKey = asset.objectKey?.trim();
      const mediaSrc = contextAssetMediaSrc(asset).trim();
      const objectKeyOption = objectKey
        ? attachmentOptions.find(
            ({ attachment }) => attachment.objectKey?.trim() === objectKey,
          )
        : undefined;
      const option =
        objectKeyOption ||
        attachmentOptions.find(({ attachment }) => {
          return (
            attachment.imageSrc.trim() === mediaSrc ||
            attachment.thumbnailSrc?.trim() === mediaSrc
          );
        });
      const matchedAttachment = option?.attachment;
      return matchedAttachment?.referenceId.trim() &&
        matchedAttachment.imageSrc.trim()
        ? appendAttachments([matchedAttachment])
        : false;
    },
    [appendAttachments, attachmentOptions],
  );

  const handleInputPaste = useCallback(
    (event: ClipboardEvent) => {
      const imageFiles = extractClipboardImageFiles(event.clipboardData);
      if (!imageFiles.length || isPastingAttachmentRef.current) {
        return;
      }

      event.preventDefault();
      const remainingAttachmentSlots =
        EDITOR_AGENT_MAX_ATTACHMENTS - attachmentsRef.current.length;
      if (remainingAttachmentSlots <= 0) {
        setAttachmentError(`最多 ${EDITOR_AGENT_MAX_ATTACHMENTS} 张`);
        return;
      }

      const uploadFiles = imageFiles.slice(0, remainingAttachmentSlots);
      const hasOverflow = uploadFiles.length < imageFiles.length;
      const currentProjectId = projectId?.trim();
      if (!currentProjectId) {
        setAttachmentError('缺少画布项目');
        return;
      }

      isPastingAttachmentRef.current = true;
      setIsPastingAttachment(true);
      setAttachmentError('图片上传中');
      void Promise.all(
        uploadFiles.map((file) =>
          createPastedAgentImageAttachment(currentProjectId, file),
        ),
      )
        .then((pastedAttachments) => {
          if (appendAttachments(pastedAttachments) && hasOverflow) {
            setAttachmentError(`最多 ${EDITOR_AGENT_MAX_ATTACHMENTS} 张`);
          }
        })
        .catch(() => {
          setAttachmentError('图片粘贴失败，请重试');
        })
        .finally(() => {
          isPastingAttachmentRef.current = false;
          setIsPastingAttachment(false);
        });
    },
    [appendAttachments, projectId],
  );

  const removeAttachment = useCallback(
    (targetAttachment: EditorAgentAttachmentRef) => {
      const key = attachmentKey(targetAttachment);
      updateAttachments((currentAttachments) =>
        currentAttachments.filter(
          (attachment) => attachmentKey(attachment) !== key,
        ),
      );
    },
    [updateAttachments],
  );

  const consumeAttachments = useCallback(() => {
    const currentAttachments = attachmentsRef.current;
    const refreshedAttachments = currentAttachments.map(
      (attachment) =>
        attachmentOptionsByKey.get(attachmentKey(attachment))?.attachment ??
        attachment,
    );
    updateAttachments([]);
    return refreshedAttachments;
  }, [attachmentOptionsByKey, updateAttachments]);

  const restoreAttachments = useCallback(
    (failedAttachments: EditorAgentAttachmentRef[]) => {
      updateAttachments((currentAttachments) => {
        // respect current (version) when duplicated with old ones
        const restoredAttachments = mergeAttachments(
          currentAttachments,
          failedAttachments,
        );
        if (restoredAttachments.length > EDITOR_AGENT_MAX_ATTACHMENTS) {
          setAttachmentError(
            `最多 ${EDITOR_AGENT_MAX_ATTACHMENTS} 张，发送失败的附件未恢复`,
          );
          return currentAttachments;
        }
        setAttachmentError(null);
        return restoredAttachments;
      });
    },
    [updateAttachments],
  );

  return {
    attachments,
    attachmentError,
    isPastingAttachment,
    attachmentPickerOpen,
    attachmentPickerTab,
    draftAttachmentKeys,
    canvasAttachmentOptions,
    libraryAttachmentOptions,
    setAttachmentPickerTab,
    openAttachmentPicker,
    closeAttachmentPicker,
    toggleAttachmentKey,
    applyAttachmentSelection,
    referenceContextAsset,
    handleInputPaste,
    removeAttachment,
    consumeAttachments,
    restoreAttachments,
  };
}