Function useModelDownload

src/hooks/useModelDownload.ts:16–123

(onComplete?: () => void)

}

export function useModelDownload(onComplete?: () => void) {
  const [downloads, setDownloads] = useState<Map<string, DownloadState>>(new Map());

  useEffect(() => {
    const setupListener = async () => {
      const unlistenDownload = await listenToDownloadProgress((progress: DownloadProgress) => {
        setDownloads((prev) => {
          const newMap = new Map(prev);
          const percentage = progress.completed && progress.total
            ? calculateProgress(progress.completed, progress.total)
            : 0;

          const existing = newMap.get(progress.model_name);
          
          newMap.set(progress.model_name, {
            status: progress.status,
            percentage,
            isDownloading: progress.status !== 'success',
            error: null,
            downloadedBytes: progress.completed ?? existing?.downloadedBytes,
            totalBytes: progress.total ?? existing?.totalBytes,
            speedMbps: existing?.speedMbps,
          });

          // Call onComplete when download finishes
          if (progress.status === 'success' && onComplete) {
            onComplete();
          }

          return newMap;
        });
      });

      const unlistenHf = await listenToHfDownloadProgress((progress) => {
        setDownloads((prev) => {
          const newMap = new Map(prev);
          const existing = newMap.get(progress.model_name);

          newMap.set(progress.model_name, {
            status: existing?.status ?? 'downloading',
            percentage: Math.round(progress.percentage),
            isDownloading: true,
            error: null,
            downloadedBytes: progress.bytes_downloaded,
            totalBytes: progress.total_bytes,
            speedMbps: progress.speed_mbps,
          });

          return newMap;
        });
      });

      return () => {
        unlistenDownload();
        unlistenHf();
      };
    };