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

/**
 * Upload an entire local directory tree to a remote peer via a `FileClient`.
 *
 * Walks the local `DirectoryHandleLike` lazily, creates remote directories
 * via `client.mkdir({ recursive: true })`, and uploads each file with
 * `client.write` (which routes inline or streams based on size). A bounded
 * concurrency pool keeps memory and inflight RPCs in check.
 *
 * Returns a `BulkTransferHandle` whose `events` stream emits `'plan'`,
 * `'file-start'`, `'file-progress'` (currently emitted at file start +
 * end), `'file-done'` / `'file-error'`, aggregate `'progress'`, and a
 * final `'complete'` (or `'abort'`).
 */
import { posixJoin } from '../utils/path.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';

interface PlannedFile {
  /** Local file handle. */
  handle: FileHandleLike;
  /** Path relative to the upload root. */
  relativePath: string;
  /** Absolute remote path. */
  remoteAbsPath: string;
}

interface PlannedDir {
  remoteAbsPath: string;
}

export interface UploadDirectoryOptions extends BulkOpts {
  /**
   * Pre-create remote directories before uploading files. Default true.
   * Disable if the server-side `write` already mkdir-p's parents.
   */
  precreateDirs?: boolean;
}

export function uploadDirectory(
  client: Pick<FileClient, 'write' | 'mkdir'>,
  local: DirectoryHandleLike,
  remoteRoot: string,
  opts: UploadDirectoryOptions = {},
): BulkTransferHandle {
  const concurrency = Math.max(
    1,
    Math.min(MAX_BULK_CONCURRENCY, opts.concurrency ?? DEFAULT_BULK_CONCURRENCY),
  );
  const precreateDirs = opts.precreateDirs ?? true;
  const continueOnError = opts.continueOnError ?? false;
  const externalSignal = opts.signal;

  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;
  let bytesTotal = 0;

  // Run the upload pipeline asynchronously.
  void (async () => {
    try {
      // 1. Plan: walk local tree, collect dirs + files
      const plannedDirs: PlannedDir[] = [];
      const plannedFiles: PlannedFile[] = [];
      try {
        await collect(local, '', remoteRoot, plannedDirs, plannedFiles, combinedSignal);
      } catch (err) {
        emit({ type: 'abort', reason: errMsg(err) });
        closeEvents();
        rejectDone(err);
        return;
      }

      const totalFiles = plannedFiles.length;
      bytesTotal = plannedFiles.reduce(
        (acc, f) => acc + (f.handle as FileHandleLike & { _size?: number })._size!,
        0,
      );
      // Note: bytesTotal is computed lazily in collect() via cached sizes.

      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 remote directories sequentially (cheap, avoids races
      //    when many uploads target the same parent).
      if (precreateDirs) {
        for (const d of plannedDirs) {
          if (combinedSignal.aborted) break;
          try {
            await client.mkdir(d.remoteAbsPath, { recursive: true });
          } catch (err) {
            if (!isAlreadyExistsError(err)) {
              if (!continueOnError) {
                emit({ type: 'abort', reason: errMsg(err) });
                closeEvents();
                rejectDone(err);
                return;
              }
            }
          }
        }
      }

      // 3. Upload files with bounded concurrency
      try {
        await runWithConcurrency(
          asyncIterableOf(plannedFiles),
          async (planned) => {
            if (combinedSignal.aborted) {
              throw new CancelledError('upload aborted');
            }
            const file = await planned.handle.getFile();
            const size = file.size;
            emit({ type: 'file-start', path: planned.relativePath, size });
            try {
              const writeOpts: { contentType?: string; signal?: AbortSignal } = {};
              if (file.type !== '') writeOpts.contentType = file.type;
              writeOpts.signal = combinedSignal;
              if (size === 0) {
                // Edge case: empty file — write empty buffer.
                await client.write(planned.remoteAbsPath, new Uint8Array(0), writeOpts);
              } else {
                await client.write(
                  planned.remoteAbsPath,
                  { stream: file.stream(), size, ...(file.type !== '' ? { contentType: file.type } : {}) },
                  writeOpts,
                );
              }
              filesDone++;
              bytesDone += size;
              emit({ type: 'file-done', path: planned.relativePath, bytesDone: size });
              emit({
                type: 'progress',
                filesDone,
                filesTotal: totalFiles,
                bytesDone,
                bytesTotal,
              });
            } catch (err) {
              emit({ type: 'file-error', path: planned.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) {
      // Belt-and-suspenders: any unhandled error reaches here.
      closeEvents();
      rejectDone(err);
    }
  })();

  // Suppress unhandled-rejection until consumer awaits done().
  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>> {
            // Caller broke out — stop iterating but keep the bulk going.
            return Promise.resolve({ value: undefined as never, done: true });
          },
        };
      },
    },
    async abort(reason) {
      internalAbort.abort(new CancelledError(reason ?? 'manual abort'));
    },
    done: () => donePromise,
  };
}

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

async function collect(
  dir: DirectoryHandleLike,
  relPrefix: string,
  remoteRoot: string,
  plannedDirs: PlannedDir[],
  plannedFiles: PlannedFile[],
  signal: AbortSignal,
): Promise<void> {
  for await (const [name, handle] of dir.entries()) {
    if (signal.aborted) return;
    const relPath = relPrefix === '' ? name : `${relPrefix}/${name}`;
    if (handle.kind === 'directory') {
      const remoteAbs = posixJoin(remoteRoot, relPath);
      plannedDirs.push({ remoteAbsPath: remoteAbs });
      await collect(handle as DirectoryHandleLike, relPath, remoteRoot, plannedDirs, plannedFiles, signal);
    } else {
      // Cache size on the handle so the planning total is accurate.
      const file = await (handle as FileHandleLike).getFile();
      (handle as FileHandleLike & { _size?: number })._size = file.size;
      plannedFiles.push({
        handle: handle as FileHandleLike,
        relativePath: relPath,
        remoteAbsPath: posixJoin(remoteRoot, relPath),
      });
    }
  }
}

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

function isAlreadyExistsError(err: unknown): boolean {
  if (err === null || typeof err !== 'object') return false;
  const code = (err as { code?: string; payload?: { code?: string } }).code;
  if (code === 'SHADE_FS_CONFLICT') return true;
  const payloadCode = (err as { payload?: { code?: string } }).payload?.code;
  return payloadCode === 'CONFLICT';
}

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