Shade/packages/shade-files/src/client/download-directory.ts

/**
 * Download an entire remote directory tree to a local `DirectoryHandleLike`.
 *
 * Mirror image of `uploadDirectory`: walks the remote tree via the shared
 * `walk()` helper, creates local directories on the fly, and downloads each
 * file with `client.read` (which routes inline or streams based on the
 * server's response). Bounded concurrency keeps RPC inflight count low.
 */
import { walk, type WalkItem } from './walk.js';
import { runWithConcurrency } from './concurrency.js';
import type { FileClient } from './client.js';
import {
  DEFAULT_BULK_CONCURRENCY,
  MAX_BULK_CONCURRENCY,
  type BulkOpts,
  type BulkTransferEvent,
  type BulkTransferHandle,
  type BulkTransferResult,
  type DirectoryHandleLike,
  type FileHandleLike,
} from './directory-types.js';
import { CancelledError } from '../schemas/errors.js';

export interface DownloadDirectoryOptions extends BulkOpts {
  /** Page size hint forwarded to the underlying `walk`. Default 200. */
  pageSize?: number;
  /** Skip files already present locally. Default false (overwrite). */
  skipExisting?: boolean;
}

export function downloadDirectory(
  client: Pick<FileClient, 'read' | 'list'>,
  remoteRoot: string,
  local: DirectoryHandleLike,
  opts: DownloadDirectoryOptions = {},
): BulkTransferHandle {
  const concurrency = Math.max(
    1,
    Math.min(MAX_BULK_CONCURRENCY, opts.concurrency ?? DEFAULT_BULK_CONCURRENCY),
  );
  const continueOnError = opts.continueOnError ?? false;
  const externalSignal = opts.signal;
  const pageSize = opts.pageSize ?? 200;

  const internalAbort = new AbortController();
  const combinedSignal = mergeSignals(externalSignal, internalAbort.signal);

  const events: BulkTransferEvent[] = [];
  const eventResolvers: ((v: IteratorResult<BulkTransferEvent>) => void)[] = [];
  let eventsClosed = false;

  function emit(event: BulkTransferEvent): void {
    if (eventsClosed) return;
    if (eventResolvers.length > 0) {
      eventResolvers.shift()!({ value: event, done: false });
    } else {
      events.push(event);
    }
  }

  function closeEvents(): void {
    if (eventsClosed) return;
    eventsClosed = true;
    while (eventResolvers.length > 0) {
      eventResolvers.shift()!({ value: undefined as never, done: true });
    }
  }

  let resolveDone!: (r: BulkTransferResult) => void;
  let rejectDone!: (err: unknown) => void;
  const donePromise = new Promise<BulkTransferResult>((resolve, reject) => {
    resolveDone = resolve;
    rejectDone = reject;
  });

  const startedAt = Date.now();
  let filesDone = 0;
  let bytesDone = 0;

  void (async () => {
    try {
      // 1. Plan: walk the remote tree into [dirs, files]
      const dirItems: WalkItem[] = [];
      const fileItems: WalkItem[] = [];
      try {
        for await (const item of walk(client, remoteRoot, {
          pageSize,
          ...(combinedSignal !== undefined ? { signal: combinedSignal } : {}),
        })) {
          if (item.entry.kind === 'dir') dirItems.push(item);
          else fileItems.push(item);
        }
      } catch (err) {
        emit({ type: 'abort', reason: errMsg(err) });
        closeEvents();
        rejectDone(err);
        return;
      }

      const totalFiles = fileItems.length;
      const bytesTotal = fileItems.reduce((acc, f) => acc + f.entry.size, 0);
      emit({ type: 'plan', totalFiles, totalBytes: bytesTotal });

      if (combinedSignal.aborted) {
        emit({ type: 'abort', reason: errMsg(combinedSignal.reason) });
        closeEvents();
        rejectDone(combinedSignal.reason ?? new CancelledError());
        return;
      }

      // 2. Pre-create local directories sequentially.
      for (const d of dirItems) {
        if (combinedSignal.aborted) break;
        try {
          await ensureLocalDir(local, d.relativePath);
        } catch (err) {
          if (!continueOnError) {
            emit({ type: 'abort', reason: errMsg(err) });
            closeEvents();
            rejectDone(err);
            return;
          }
        }
      }

      // 3. Download files with bounded concurrency.
      try {
        await runWithConcurrency(
          asyncIterableOf(fileItems),
          async (item) => {
            if (combinedSignal.aborted) {
              throw new CancelledError('download aborted');
            }
            emit({ type: 'file-start', path: item.relativePath, size: item.entry.size });
            try {
              const fileHandle = await ensureLocalFile(local, item.relativePath);
              if (opts.skipExisting === true) {
                // Skip if file already exists with non-zero size
                try {
                  const existing = await fileHandle.getFile();
                  if (existing.size === item.entry.size) {
                    filesDone++;
                    bytesDone += existing.size;
                    emit({ type: 'file-done', path: item.relativePath, bytesDone: existing.size });
                    emit({
                      type: 'progress',
                      filesDone,
                      filesTotal: totalFiles,
                      bytesDone,
                      bytesTotal,
                    });
                    return;
                  }
                } catch {
                  /* not present yet — fall through to download */
                }
              }
              const writable = await fileHandle.createWritable();
              try {
                const result = await client.read(item.absolutePath, {
                  signal: combinedSignal,
                });
                if (result.kind === 'inline') {
                  if (result.bytes.byteLength > 0) await writable.write(result.bytes);
                  await writable.close();
                  filesDone++;
                  bytesDone += result.bytes.byteLength;
                  emit({ type: 'file-done', path: item.relativePath, bytesDone: result.bytes.byteLength });
                } else {
                  await pipeReadableToWritable(result.stream, writable);
                  await writable.close();
                  await result.done();
                  filesDone++;
                  bytesDone += result.size;
                  emit({ type: 'file-done', path: item.relativePath, bytesDone: result.size });
                }
                emit({
                  type: 'progress',
                  filesDone,
                  filesTotal: totalFiles,
                  bytesDone,
                  bytesTotal,
                });
              } catch (err) {
                await writable.abort(err).catch(() => undefined);
                throw err;
              }
            } catch (err) {
              emit({ type: 'file-error', path: item.relativePath, error: err });
              throw err;
            }
          },
          {
            concurrency,
            signal: combinedSignal,
            continueOnError,
            onError: () => {
              /* already emitted as 'file-error' */
            },
          },
        );
      } catch (err) {
        emit({ type: 'abort', reason: errMsg(err) });
        closeEvents();
        rejectDone(err);
        return;
      }

      const durationMs = Date.now() - startedAt;
      emit({ type: 'complete', filesDone, bytesDone, durationMs });
      closeEvents();
      resolveDone({ filesDone, bytesDone, durationMs });
    } catch (err) {
      closeEvents();
      rejectDone(err);
    }
  })();

  donePromise.catch(() => {
    /* deliberate */
  });

  return {
    events: {
      [Symbol.asyncIterator]() {
        return {
          next(): Promise<IteratorResult<BulkTransferEvent>> {
            if (events.length > 0) {
              return Promise.resolve({ value: events.shift()!, done: false });
            }
            if (eventsClosed) {
              return Promise.resolve({ value: undefined as never, done: true });
            }
            return new Promise((resolve) => eventResolvers.push(resolve));
          },
          return(): Promise<IteratorResult<BulkTransferEvent>> {
            return Promise.resolve({ value: undefined as never, done: true });
          },
        };
      },
    },
    async abort(reason) {
      internalAbort.abort(new CancelledError(reason ?? 'manual abort'));
    },
    done: () => donePromise,
  };
}

// ─── Helpers ─────────────────────────────────────────────────

async function ensureLocalDir(
  root: DirectoryHandleLike,
  relativePath: string,
): Promise<DirectoryHandleLike> {
  const segments = relativePath.split('/').filter((s) => s !== '');
  let current = root;
  for (const seg of segments) {
    current = await current.getDirectoryHandle(seg, { create: true });
  }
  return current;
}

async function ensureLocalFile(
  root: DirectoryHandleLike,
  relativePath: string,
): Promise<FileHandleLike> {
  const segments = relativePath.split('/').filter((s) => s !== '');
  if (segments.length === 0) throw new Error('empty file path');
  let current = root;
  for (let i = 0; i < segments.length - 1; i++) {
    current = await current.getDirectoryHandle(segments[i]!, { create: true });
  }
  return await current.getFileHandle(segments[segments.length - 1]!, { create: true });
}

async function pipeReadableToWritable(
  readable: ReadableStream<Uint8Array>,
  writable: { write(chunk: Uint8Array): Promise<void> },
): Promise<void> {
  const reader = readable.getReader();
  try {
    while (true) {
      const { value, done } = await reader.read();
      if (done) break;
      if (value !== undefined && value.byteLength > 0) {
        await writable.write(value);
      }
    }
  } finally {
    reader.releaseLock();
  }
}

async function* asyncIterableOf<T>(arr: T[]): AsyncIterable<T> {
  for (const item of arr) yield item;
}

function errMsg(err: unknown): string {
  if (err instanceof Error) return err.message;
  return String(err ?? 'unknown error');
}

function mergeSignals(
  external: AbortSignal | undefined,
  internal: AbortSignal,
): AbortSignal {
  if (external === undefined) return internal;
  if (external.aborted) return external;
  const controller = new AbortController();
  const onAbort = (sig: AbortSignal): void => {
    controller.abort(sig.reason);
  };
  external.addEventListener('abort', () => onAbort(external), { once: true });
  internal.addEventListener('abort', () => onAbort(internal), { once: true });
  return controller.signal;
}