DataMate/frontend/src/hooks/useSliceUpload.tsx

import { TaskItem } from "@/pages/DataManagement/dataset.model";
import { calculateSHA256, checkIsFilesExist, streamSplitAndUpload, StreamUploadResult } from "@/utils/file.util";
import { App } from "antd";
import { useRef, useState } from "react";

export function useFileSliceUpload(
  {
    preUpload,
    uploadChunk,
    cancelUpload,
  }: {
    preUpload: (id: string, params: Record<string, unknown>) => Promise<{ data: number }>;
    uploadChunk: (id: string, formData: FormData, config: Record<string, unknown>) => Promise<unknown>;
    cancelUpload: ((reqId: number) => Promise<unknown>) | null;
  },
  showTaskCenter = true, // 上传时是否显示任务中心
  enableStreamUpload = true // 是否启用流式分割上传
) {
  const { message } = App.useApp();
  const [taskList, setTaskList] = useState<TaskItem[]>([]);
  const taskListRef = useRef<TaskItem[]>([]); // 用于固定任务顺序

  const createTask = (detail: Record<string, unknown> = {}) => {
    const { dataset } = detail;
    const title = `上传数据集: ${dataset.name} `;
    const controller = new AbortController();
    const task: TaskItem = {
      key: dataset.id,
      title,
      percent: 0,
      reqId: -1,
      controller,
      size: 0,
      updateEvent: detail.updateEvent,
      hasArchive: detail.hasArchive,
      prefix: detail.prefix,
    };
    taskListRef.current = [task, ...taskListRef.current];

    setTaskList(taskListRef.current);
    
    // 立即显示任务中心，让用户感知上传已开始
    if (showTaskCenter) {
      window.dispatchEvent(
        new CustomEvent("show:task-popover", { detail: { show: true } })
      );
    }
    
    return task;
  };

  const updateTaskList = (task: TaskItem) => {
    taskListRef.current = taskListRef.current.map((item) =>
      item.key === task.key ? task : item
    );
    setTaskList(taskListRef.current);
  };

  const removeTask = (task: TaskItem) => {
    const { key } = task;
    taskListRef.current = taskListRef.current.filter(
      (item) => item.key !== key
    );
    setTaskList(taskListRef.current);
    if (task.isCancel && task.cancelFn) {
      task.cancelFn();
    }
    if (task.updateEvent) {
      // 携带前缀信息，便于刷新后仍停留在当前目录
      window.dispatchEvent(
        new CustomEvent(task.updateEvent, {
          detail: { prefix: task.prefix },
        })
      );
    }
    if (showTaskCenter) {
      window.dispatchEvent(
        new CustomEvent("show:task-popover", { detail: { show: false } })
      );
    }
  };

  async function buildFormData({ file, reqId, i, j }: { file: { slices: Blob[]; name: string; size: number }; reqId: number; i: number; j: number }) {
    const formData = new FormData();
    const { slices, name, size } = file;
    const checkSum = await calculateSHA256(slices[j]);
    formData.append("file", slices[j]);
    formData.append("reqId", reqId.toString());
    formData.append("fileNo", (i + 1).toString());
    formData.append("chunkNo", (j + 1).toString());
    formData.append("fileName", name);
    formData.append("fileSize", size.toString());
    formData.append("totalChunkNum", slices.length.toString());
    formData.append("checkSumHex", checkSum);
    return formData;
  }

  async function uploadSlice(task: TaskItem, fileInfo: { loaded: number; i: number; j: number; files: { slices: Blob[]; name: string; size: number }[]; totalSize: number }) {
    if (!task) {
      return;
    }
    const { reqId, key, controller } = task;
    const { loaded, i, j, files, totalSize } = fileInfo;
    
    // 检查是否已取消
    if (controller.signal.aborted) {
      throw new Error("Upload cancelled");
    }
    
    const formData = await buildFormData({
      file: files[i],
      i,
      j,
      reqId,
    });

    let newTask = { ...task };
    await uploadChunk(key, formData, {
      signal: controller.signal,
      onUploadProgress: (e) => {
        const loadedSize = loaded + e.loaded;
        const curPercent = Number((loadedSize / totalSize) * 100).toFixed(2);

        newTask = {
          ...newTask,
          ...taskListRef.current.find((item) => item.key === key),
          size: loadedSize,
          percent: curPercent >= 100 ? 99.99 : curPercent,
        };
        updateTaskList(newTask);
      },
    });
  }

  async function uploadFile({ task, files, totalSize }: { task: TaskItem; files: { slices: Blob[]; name: string; size: number; originFile: Blob }[]; totalSize: number }) {
    console.log('[useSliceUpload] Calling preUpload with prefix:', task.prefix);
    const { data: reqId } = await preUpload(task.key, {
      totalFileNum: files.length,
      totalSize,
      datasetId: task.key,
      hasArchive: task.hasArchive,
      prefix: task.prefix,
    });
    console.log('[useSliceUpload] PreUpload response reqId:', reqId);

    const newTask: TaskItem = {
      ...task,
      reqId,
      isCancel: false,
      cancelFn: () => {
        // 使用 newTask 的 controller 确保一致性
        newTask.controller.abort();
        cancelUpload?.(reqId);
        if (newTask.updateEvent) window.dispatchEvent(new Event(newTask.updateEvent));
      },
    };
    updateTaskList(newTask);
    // 注意：show:task-popover 事件已在 createTask 中触发，此处不再重复触发
    // // 更新数据状态
    if (task.updateEvent) window.dispatchEvent(new Event(task.updateEvent));

    let loaded = 0;
    for (let i = 0; i < files.length; i++) {
      // 检查是否已取消
      if (newTask.controller.signal.aborted) {
        throw new Error("Upload cancelled");
      }
      const { slices } = files[i];
      for (let j = 0; j < slices.length; j++) {
        // 检查是否已取消
        if (newTask.controller.signal.aborted) {
          throw new Error("Upload cancelled");
        }
        await uploadSlice(newTask, {
          loaded,
          i,
          j,
          files,
          totalSize,
        });
        loaded += slices[j].size;
      }
    }
    removeTask(newTask);
  }

  const handleUpload = async ({ task, files }: { task: TaskItem; files: { slices: Blob[]; name: string; size: number; originFile: Blob }[] }) => {
    const isErrorFile = await checkIsFilesExist(files);
    if (isErrorFile) {
      message.error("文件被修改或删除，请重新选择文件上传");
      removeTask({
        ...task,
        isCancel: false,
        ...taskListRef.current.find((item) => item.key === task.key),
      });
      return;
    }

    try {
      const totalSize = files.reduce((acc, file) => acc + file.size, 0);
      await uploadFile({ task, files, totalSize });
    } catch (err) {
      console.error(err);
      message.error("文件上传失败，请稍后重试");
      removeTask({
        ...task,
        isCancel: true,
        ...taskListRef.current.find((item) => item.key === task.key),
      });
    }
  };

  /**
   * 流式分割上传处理
   * 用于大文件按行分割并立即上传的场景
   */
  const handleStreamUpload = async ({ task, files }: { task: TaskItem; files: File[] }) => {
    try {
      console.log('[useSliceUpload] Starting stream upload for', files.length, 'files');
      
      const totalSize = files.reduce((acc, file) => acc + file.size, 0);
      
      // 存储所有文件的 reqId，用于取消上传
      const reqIds: number[] = [];
      
      const newTask: TaskItem = {
        ...task,
        reqId: -1,
        isCancel: false,
        cancelFn: () => {
          // 使用 newTask 的 controller 确保一致性
          newTask.controller.abort();
          // 取消所有文件的预上传请求
          reqIds.forEach(id => cancelUpload?.(id));
          if (newTask.updateEvent) window.dispatchEvent(new Event(newTask.updateEvent));
        },
      };
      updateTaskList(newTask);

      let totalUploadedLines = 0;
      let totalProcessedBytes = 0;
      const results: StreamUploadResult[] = [];

      // 逐个处理文件，每个文件单独调用 preUpload
      for (let i = 0; i < files.length; i++) {
        // 检查是否已取消
        if (newTask.controller.signal.aborted) {
          throw new Error("Upload cancelled");
        }
        
        const file = files[i];
        console.log(`[useSliceUpload] Processing file ${i + 1}/${files.length}: ${file.name}`);
        
        const result = await streamSplitAndUpload(
          file,
          (formData, config) => uploadChunk(task.key, formData, {
            ...config,
            signal: newTask.controller.signal,
          }),
          (currentBytes, totalBytes, uploadedLines) => {
            // 检查是否已取消
            if (newTask.controller.signal.aborted) {
              return;
            }
            
            // 更新进度
            const overallBytes = totalProcessedBytes + currentBytes;
            const curPercent = Number((overallBytes / totalSize) * 100).toFixed(2);
            
            const updatedTask: TaskItem = {
              ...newTask,
              ...taskListRef.current.find((item) => item.key === task.key),
              size: overallBytes,
              percent: curPercent >= 100 ? 99.99 : curPercent,
              streamUploadInfo: {
                currentFile: file.name,
                fileIndex: i + 1,
                totalFiles: files.length,
                uploadedLines: totalUploadedLines + uploadedLines,
              },
            };
            updateTaskList(updatedTask);
          }, 
          1024 * 1024, // 1MB chunk size
          {
            resolveReqId: async ({ totalFileNum, totalSize }) => {
              const { data: reqId } = await preUpload(task.key, {
                totalFileNum,
                totalSize,
                datasetId: task.key,
                hasArchive: task.hasArchive,
                prefix: task.prefix,
              });
              console.log(`[useSliceUpload] File ${file.name} preUpload response reqId:`, reqId);
              reqIds.push(reqId);
              return reqId;
            },
            hasArchive: newTask.hasArchive,
            prefix: newTask.prefix,
            signal: newTask.controller.signal,
            maxConcurrency: 3,
          }
        );
        
        results.push(result);
        totalUploadedLines += result.uploadedCount;
        totalProcessedBytes += file.size;
        
        console.log(`[useSliceUpload] File ${file.name} processed, uploaded ${result.uploadedCount} lines`);
      }

      console.log('[useSliceUpload] Stream upload completed, total lines:', totalUploadedLines);
      removeTask(newTask);
      
      message.success(`成功上传 ${totalUploadedLines} 个文件（按行分割）`);
    } catch (err) {
      console.error('[useSliceUpload] Stream upload error:', err);
      if (err.message === "Upload cancelled") {
        message.info("上传已取消");
      } else {
        message.error("文件上传失败，请稍后重试");
      }
      removeTask({
        ...task,
        isCancel: true,
        ...taskListRef.current.find((item) => item.key === task.key),
      });
    }
  };

  /**
   * 注册流式上传事件监听
   * 返回注销函数
   */
  const registerStreamUploadListener = () => {
    if (!enableStreamUpload) return () => {};

    const streamUploadHandler = async (e: Event) => {
      const customEvent = e as CustomEvent;
      const { dataset, files, updateEvent, hasArchive, prefix } = customEvent.detail;
      
      const controller = new AbortController();
      const task: TaskItem = {
        key: dataset.id,
        title: `上传数据集: ${dataset.name} (按行分割)`,
        percent: 0,
        reqId: -1,
        controller,
        size: 0,
        updateEvent,
        hasArchive,
        prefix,
      };
      
      taskListRef.current = [task, ...taskListRef.current];
      setTaskList(taskListRef.current);
      
      // 显示任务中心
      if (showTaskCenter) {
        window.dispatchEvent(
          new CustomEvent("show:task-popover", { detail: { show: true } })
        );
      }
      
      await handleStreamUpload({ task, files });
    };

    window.addEventListener("upload:dataset-stream", streamUploadHandler);
    
    return () => {
      window.removeEventListener("upload:dataset-stream", streamUploadHandler);
    };
  };

  return {
    taskList,
    createTask,
    removeTask,
    handleUpload,
    handleStreamUpload,
    registerStreamUploadListener,
  };
}