|
|
|
|
|
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 |
|
|
}; |
|
|
}; |
|
|
|