EasyFlow/easyflow-ui-admin/packages/effects/common-ui/src/components/chat-timeline/useChatDocumentUploads.ts

import type { ChatDocumentAttachment } from './types';

import { computed, ref } from 'vue';

import { uuid } from '@easyflow/utils';

const MAX_DOCUMENTS = 3;
const MAX_TOTAL_BYTES = 30 * 1024 * 1024;
const MAX_OFFICE_BYTES = 20 * 1024 * 1024;
const MAX_EXCEL_BYTES = 10 * 1024 * 1024;
const MAX_TEXT_BYTES = 5 * 1024 * 1024;
const POLL_INTERVAL_MS = 750;
const MAX_POLL_ATTEMPTS = 80;
export const CHAT_DOCUMENT_SUPPORTED_FORMATS =
  'PDF、Word、PPT、Excel、TXT、Markdown';
const SUPPORTED_EXTENSIONS = new Set([
  'doc',
  'docx',
  'md',
  'pdf',
  'ppt',
  'pptx',
  'txt',
  'xls',
  'xlsx',
]);
const EXCEL_EXTENSIONS = new Set(['xls', 'xlsx']);
const TEXT_EXTENSIONS = new Set(['md', 'txt']);

export interface ChatDocumentUploadContext {
  agentId: string;
  mode: 'DRAFT' | 'FORMAL';
  sessionId: string;
}

export interface ChatDocumentUploadView extends ChatDocumentAttachment {
  expiresAt?: string;
  status: 'error' | 'reading' | 'ready' | 'uploading';
  uploadId: string;
}

interface RequestResult<T = any> {
  data: T;
  errorCode: number;
  message?: string;
}

export interface ChatDocumentUploadApi {
  delete: (uploadId: string) => Promise<RequestResult<void>>;
  retry: (uploadId: string) => Promise<RequestResult<ChatDocumentUploadView>>;
  status: (uploadId: string) => Promise<RequestResult<ChatDocumentUploadView>>;
  upload: (
    file: File,
    context: ChatDocumentUploadContext,
    uploadId: string,
  ) => Promise<RequestResult<ChatDocumentUploadView>>;
}

export interface ChatDocumentAddFilesResult {
  rejectedCount: number;
  unsupportedCount: number;
}

interface LocalDocument extends ChatDocumentAttachment {
  file?: File;
}

function createLocalId() {
  return `agent-document-${Date.now()}-${Math.random().toString(16).slice(2, 8)}`;
}

function createUploadId() {
  return uuid().replaceAll('-', '');
}

function extensionOf(file: File) {
  return file.name.includes('.')
    ? file.name.split('.').pop()?.toLowerCase() || ''
    : '';
}

function maxBytes(extension: string) {
  if (EXCEL_EXTENSIONS.has(extension)) return MAX_EXCEL_BYTES;
  if (TEXT_EXTENSIONS.has(extension)) return MAX_TEXT_BYTES;
  return MAX_OFFICE_BYTES;
}

function sizeLimitMessage(extension: string) {
  if (EXCEL_EXTENSIONS.has(extension)) {
    return 'Excel 文档不能超过 10 MiB';
  }
  if (TEXT_EXTENSIONS.has(extension)) {
    return '文本文件不能超过 5 MiB';
  }
  return '单份文档不能超过 20 MiB';
}

function errorMessage(error: unknown) {
  const candidate = error as any;
  return (
    candidate?.response?.data?.message || candidate?.message || '文档上传失败'
  );
}

function displayStatus(
  value?: string,
): NonNullable<ChatDocumentAttachment['status']> {
  const status = String(value || '').toUpperCase();
  if (status === 'READY') return 'ready';
  if (status === 'FAILED' || status === 'READ_FAILED' || status === 'EXPIRED') {
    return 'error';
  }
  if (status === 'UPLOADING') return 'uploading';
  return 'reading';
}

function applyServerView(target: LocalDocument, view: ChatDocumentUploadView) {
  const status = displayStatus(view.status);
  Object.assign(target, view, {
    error: view.error || (view as any).errorMessage,
    file: target.file,
    localId: target.localId,
    status,
  });
  return status;
}

function wait(milliseconds: number) {
  return new Promise((resolve) => setTimeout(resolve, milliseconds));
}

export function createChatDocumentUploads(api: ChatDocumentUploadApi) {
  const items = ref<LocalDocument[]>([]);
  let lifecycle = 0;

  const readyItems = computed(() =>
    items.value.filter(
      (item): item is ChatDocumentUploadView & LocalDocument =>
        item.status === 'ready' && Boolean(item.uploadId),
    ),
  );
  const uploadIds = computed(() =>
    readyItems.value.map((item) => item.uploadId),
  );
  const processing = computed(() =>
    items.value.some(
      (item) => item.status === 'uploading' || item.status === 'reading',
    ),
  );

  async function poll(item: LocalDocument, expectedLifecycle = lifecycle) {
    if (!item.uploadId) return;
    for (let attempt = 0; attempt < MAX_POLL_ATTEMPTS; attempt++) {
      await wait(POLL_INTERVAL_MS);
      if (expectedLifecycle !== lifecycle || !items.value.includes(item)) {
        return;
      }
      try {
        const response = await api.status(item.uploadId);
        if (response.errorCode !== 0 || !response.data) {
          throw new Error(response.message || '文档读取状态查询失败');
        }
        applyServerView(item, response.data);
        if (item.status === 'ready' || item.status === 'error') {
          return;
        }
      } catch (error) {
        item.status = 'error';
        item.error = errorMessage(error);
        return;
      }
    }
    item.status = 'error';
    item.error = '文档读取时间较长，请重试状态';
  }

  async function upload(
    item: LocalDocument,
    context: ChatDocumentUploadContext,
  ) {
    if (!item.file) return;
    item.status = 'uploading';
    item.error = undefined;
    const currentLifecycle = lifecycle;
    try {
      item.uploadId ||= createUploadId();
      const response = await api.upload(item.file, context, item.uploadId);
      if (response.errorCode !== 0 || !response.data) {
        throw new Error(response.message || '文档上传失败');
      }
      if (currentLifecycle !== lifecycle || !items.value.includes(item)) {
        await api.delete(response.data.uploadId);
        return;
      }
      const status = applyServerView(item, response.data);
      if (status !== 'ready' && status !== 'error') {
        await poll(item, currentLifecycle);
      }
    } catch (error) {
      if (items.value.includes(item)) {
        item.status = 'error';
        item.error = errorMessage(error);
      }
    }
  }

  async function addFiles(files: File[], context: ChatDocumentUploadContext) {
    const supportedFiles: File[] = [];
    let unsupportedCount = 0;
    for (const file of files) {
      if (SUPPORTED_EXTENSIONS.has(extensionOf(file))) {
        supportedFiles.push(file);
      } else {
        unsupportedCount++;
      }
    }
    const available = Math.max(0, MAX_DOCUMENTS - items.value.length);
    const accepted = supportedFiles.slice(0, available);
    const rejectedCount = Math.max(0, supportedFiles.length - accepted.length);
    const currentTotal = items.value.reduce(
      (total, item) => total + Number(item.size || 0),
      0,
    );
    let addedBytes = 0;
    const uploads: Promise<void>[] = [];
    for (const file of accepted) {
      const extension = extensionOf(file);
      const item: LocalDocument = {
        file,
        localId: createLocalId(),
        mimeType: file.type,
        name: file.name || '文档',
        size: file.size,
        status: 'uploading',
      };
      items.value.push(item);
      const trackedItem = items.value[items.value.length - 1];
      if (!trackedItem) continue;
      if (file.size > maxBytes(extension)) {
        trackedItem.status = 'error';
        trackedItem.error = sizeLimitMessage(extension);
        continue;
      }
      if (currentTotal + addedBytes + file.size > MAX_TOTAL_BYTES) {
        trackedItem.status = 'error';
        trackedItem.error = '本轮文档总大小不能超过 30 MiB';
        continue;
      }
      addedBytes += file.size;
      uploads.push(upload(trackedItem, context));
    }
    await Promise.all(uploads);
    return {
      rejectedCount,
      unsupportedCount,
    } satisfies ChatDocumentAddFilesResult;
  }

  async function remove(item: ChatDocumentAttachment) {
    const index = items.value.findIndex(
      (candidate) =>
        candidate.localId === item.localId ||
        (candidate.uploadId && candidate.uploadId === item.uploadId),
    );
    if (index === -1) return;
    const selected = items.value[index];
    if (selected?.uploadId) {
      const response = await api.delete(selected.uploadId);
      if (response.errorCode !== 0) {
        throw new Error(response.message || '文档删除失败');
      }
    }
    items.value.splice(index, 1);
  }

  async function retry(
    item: ChatDocumentAttachment,
    context: ChatDocumentUploadContext,
  ) {
    const found = items.value.find(
      (candidate) =>
        candidate.localId === item.localId ||
        (candidate.uploadId && candidate.uploadId === item.uploadId),
    );
    if (!found) return;
    found.error = undefined;
    if (!found.attachmentRef && found.file) {
      await upload(found, context);
      return;
    }
    if (!found.uploadId && found.file) {
      await upload(found, context);
      return;
    }
    if (!found.uploadId) return;
    try {
      const statusResponse = await api.status(found.uploadId);
      if (statusResponse.errorCode !== 0 || !statusResponse.data) {
        throw new Error(statusResponse.message || '文档状态查询失败');
      }
      const status = applyServerView(found, statusResponse.data);
      if (status === 'ready') return;
      if (status === 'error') {
        const retryResponse = await api.retry(found.uploadId);
        if (retryResponse.errorCode !== 0 || !retryResponse.data) {
          throw new Error(retryResponse.message || '文档读取重试失败');
        }
        applyServerView(found, retryResponse.data);
      }
      if ((found.status as ChatDocumentAttachment['status']) !== 'ready') {
        await poll(found);
      }
    } catch (error) {
      found.status = 'error';
      found.error = errorMessage(error);
    }
  }

  function restore(restored: ChatDocumentUploadView[] = []) {
    clear();
    const currentLifecycle = lifecycle;
    items.value = restored.slice(0, MAX_DOCUMENTS).map((item) => ({
      ...item,
      status: displayStatus(item.status),
    }));
    for (const item of items.value) {
      if (item.status === 'reading' || item.status === 'uploading') {
        void poll(item, currentLifecycle);
      }
    }
  }

  function clear() {
    lifecycle++;
    items.value = [];
  }

  return {
    addFiles,
    clear,
    items,
    processing,
    readyItems,
    remove,
    restore,
    retry,
    uploadIds,
  };
}