import type { ImageFile } from '@/types/app' import { toast } from '@langgenius/dify-ui/toast' import { useCallback, useMemo, useRef, useState } from 'react' import { useTranslation } from 'react-i18next' import { useParams } from '@/next/navigation' import { ALLOW_FILE_EXTENSIONS, TransferMethod } from '@/types/app' import { getImageUploadErrorMessage, imageUpload } from './utils' export const useImageFiles = () => { const params = useParams() const { t } = useTranslation() const [files, setFiles] = useState([]) const filesRef = useRef([]) const handleUpload = (imageFile: ImageFile) => { const files = filesRef.current const index = files.findIndex((file) => file._id === imageFile._id) if (index > -1) { const currentFile = files[index] const newFiles = [ ...files.slice(0, index), { ...currentFile, ...imageFile }, ...files.slice(index + 1), ] setFiles(newFiles) filesRef.current = newFiles } else { const newFiles = [...files, imageFile] setFiles(newFiles) filesRef.current = newFiles } } const handleRemove = (imageFileId: string) => { const files = filesRef.current const index = files.findIndex((file) => file._id === imageFileId) if (index > -1) { const currentFile = files[index]! const newFiles = [ ...files.slice(0, index), { ...currentFile, deleted: true }, ...files.slice(index + 1), ] setFiles(newFiles) filesRef.current = newFiles } } const handleImageLinkLoadError = (imageFileId: string) => { const files = filesRef.current const index = files.findIndex((file) => file._id === imageFileId) if (index > -1) { const currentFile = files[index]! const newFiles = [ ...files.slice(0, index), { ...currentFile, progress: -1 }, ...files.slice(index + 1), ] filesRef.current = newFiles setFiles(newFiles) } } const handleImageLinkLoadSuccess = (imageFileId: string) => { const files = filesRef.current const index = files.findIndex((file) => file._id === imageFileId) if (index > -1) { const currentImageFile = files[index]! const newFiles = [ ...files.slice(0, index), { ...currentImageFile, progress: 100 }, ...files.slice(index + 1), ] filesRef.current = newFiles setFiles(newFiles) } } const handleReUpload = (imageFileId: string) => { const files = filesRef.current const index = files.findIndex((file) => file._id === imageFileId) if (index > -1) { const currentImageFile = files[index]! imageUpload( { file: currentImageFile!.file!, onProgressCallback: (progress) => { const newFiles = [ ...files.slice(0, index), { ...currentImageFile, progress }, ...files.slice(index + 1), ] filesRef.current = newFiles setFiles(newFiles) }, onSuccessCallback: (res) => { const newFiles = [ ...files.slice(0, index), { ...currentImageFile, fileId: res.id, progress: 100 }, ...files.slice(index + 1), ] filesRef.current = newFiles setFiles(newFiles) }, onErrorCallback: (error?: any) => { const errorMessage = getImageUploadErrorMessage( error, t(($) => $['imageUploader.uploadFromComputerUploadError'], { ns: 'common' }), t, ) toast.error(errorMessage) const newFiles = [ ...files.slice(0, index), { ...currentImageFile, progress: -1 }, ...files.slice(index + 1), ] filesRef.current = newFiles setFiles(newFiles) }, }, !!params?.token, ) } } const handleClear = () => { setFiles([]) filesRef.current = [] } const filteredFiles = useMemo(() => { return files.filter((file) => !file.deleted) }, [files]) return { files: filteredFiles, onUpload: handleUpload, onRemove: handleRemove, onImageLinkLoadError: handleImageLinkLoadError, onImageLinkLoadSuccess: handleImageLinkLoadSuccess, onReUpload: handleReUpload, onClear: handleClear, } } type useLocalUploaderProps = { disabled?: boolean limit?: number onUpload: (imageFile: ImageFile) => void } export const useLocalFileUploader = ({ limit, disabled = false, onUpload, }: useLocalUploaderProps) => { const params = useParams() const { t } = useTranslation() const handleLocalFileUpload = useCallback( (file: File) => { if (disabled) { // TODO: leave some warnings? return } if (!ALLOW_FILE_EXTENSIONS.includes(file.type.split('/')[1]!)) return if (limit && file.size > limit * 1024 * 1024) { toast.error( t(($) => $['imageUploader.uploadFromComputerLimit'], { ns: 'common', size: limit }), ) return } const reader = new FileReader() reader.addEventListener( 'load', () => { const imageFile = { type: TransferMethod.local_file, _id: `${Date.now()}`, fileId: '', file, url: reader.result as string, base64Url: reader.result as string, progress: 0, } onUpload(imageFile) imageUpload( { file: imageFile.file, onProgressCallback: (progress) => { onUpload({ ...imageFile, progress }) }, onSuccessCallback: (res) => { onUpload({ ...imageFile, fileId: res.id, progress: 100 }) }, onErrorCallback: (error?: any) => { const errorMessage = getImageUploadErrorMessage( error, t(($) => $['imageUploader.uploadFromComputerUploadError'], { ns: 'common' }), t, ) toast.error(errorMessage) onUpload({ ...imageFile, progress: -1 }) }, }, !!params?.token, ) }, false, ) reader.addEventListener( 'error', () => { toast.error(t(($) => $['imageUploader.uploadFromComputerReadError'], { ns: 'common' })) }, false, ) reader.readAsDataURL(file) }, [disabled, limit, t, onUpload, params?.token], ) return { disabled, handleLocalFileUpload } }