|
|
| import { useState, useCallback, useRef } from 'react'; |
| import { FileItem, UploadStatus } from '../types'; |
| import { uploadBatchToHub } from '../services/hfService'; |
|
|
| export const useFileUpload = () => { |
| const [files, setFiles] = useState<FileItem[]>([]); |
| const [isUploading, setIsUploading] = useState(false); |
| |
| |
| const filesRef = useRef<FileItem[]>([]); |
| filesRef.current = files; |
|
|
| |
| |
| const BATCH_SIZE = 10; |
| |
| const CONCURRENCY_LIMIT = 5; |
|
|
| |
|
|
| const addFiles = useCallback((newFilesList: FileItem[]) => { |
| setFiles((prev) => [...prev, ...newFilesList]); |
| }, []); |
|
|
| const removeFile = useCallback((id: string) => { |
| setFiles((prev) => prev.filter((f) => f.id !== id)); |
| }, []); |
|
|
| const updateFilePath = useCallback((id: string, newPath: string) => { |
| setFiles((prev) => prev.map((f) => (f.id === id ? { ...f, path: newPath } : f))); |
| }, []); |
|
|
| const startUpload = useCallback(async () => { |
| |
| const pendingFiles = filesRef.current.filter( |
| (f) => f.status === UploadStatus.IDLE || f.status === UploadStatus.ERROR |
| ); |
|
|
| if (pendingFiles.length === 0) return; |
|
|
| setIsUploading(true); |
|
|
| |
| const batches: FileItem[][] = []; |
| for (let i = 0; i < pendingFiles.length; i += BATCH_SIZE) { |
| batches.push(pendingFiles.slice(i, i + BATCH_SIZE)); |
| } |
|
|
| |
| const queue = [...batches]; |
| let activeWorkers = 0; |
|
|
| |
| const updateBatchStatus = (batchItems: FileItem[], status: UploadStatus, result?: { urls?: string[], error?: string }) => { |
| setFiles((prev) => |
| prev.map((f) => { |
| const batchIndex = batchItems.findIndex(b => b.id === f.id); |
| if (batchIndex !== -1) { |
| return { |
| ...f, |
| status: status, |
| url: status === UploadStatus.SUCCESS ? result?.urls?.[batchIndex] : f.url, |
| error: status === UploadStatus.ERROR ? result?.error : undefined |
| }; |
| } |
| return f; |
| }) |
| ); |
| }; |
|
|
| |
| |
| const processNextBatch = async (): Promise<void> => { |
| if (queue.length === 0) return; |
|
|
| const batch = queue.shift(); |
| if (!batch) return; |
|
|
| activeWorkers++; |
| |
| |
| updateBatchStatus(batch, UploadStatus.UPLOADING); |
|
|
| try { |
| const payload = batch.map(item => ({ |
| id: item.id, |
| file: item.file, |
| path: item.path |
| })); |
|
|
| const urls = await uploadBatchToHub(payload); |
| |
| |
| updateBatchStatus(batch, UploadStatus.SUCCESS, { urls }); |
|
|
| } catch (err: any) { |
| console.error("Batch failed:", err); |
| |
| updateBatchStatus(batch, UploadStatus.ERROR, { error: err.message || "Upload failed" }); |
| } finally { |
| activeWorkers--; |
| |
| if (queue.length > 0) { |
| await processNextBatch(); |
| } |
| } |
| }; |
|
|
| |
| |
| const initialWorkers = []; |
| const limit = Math.min(CONCURRENCY_LIMIT, batches.length); |
| |
| for (let i = 0; i < limit; i++) { |
| initialWorkers.push(processNextBatch()); |
| } |
|
|
| |
| await Promise.all(initialWorkers); |
| |
| setIsUploading(false); |
|
|
| }, []); |
|
|
| return { |
| files, |
| isUploading, |
| addFiles, |
| removeFile, |
| updateFilePath, |
| startUpload |
| }; |
| }; |
|
|