import { StreamRouteUnknownKeyError } from "../../errors/stream-error.js";
import type { RouteOptions } from "./route.types.js";

/**
 * Internal: build a `Map<K, ReadableStream<T>>` of routed branches, single-pass
 * (no chunk is duplicated). The reader drives the source and dispatches each
 * chunk to the controller of exactly one branch — chosen by `keyFn`. On
 * unknown key under `onUnknownKey: 'error'`, every branch errors with
 * `StreamRouteUnknownKeyError` and the source is cancelled.
 */
function buildRoutedMap<T, K>(
    source: ReadableStream<T>,
    keyFn: (chunk: T) => K,
    keys: readonly K[],
    options: RouteOptions = {}
): Map<K, ReadableStream<T>> {
    const onUnknownKey = options.onUnknownKey ?? "drop";

    const controllers = new Map<K, ReadableStreamDefaultController<T>>();
    const knownKeys = new Set<K>(keys);
    const result = new Map<K, ReadableStream<T>>();

    // We need every branch's controller wired up before the reader starts so
    // that early chunks are not lost when the source resolves quickly.
    const controllerReady: Promise<void>[] = [];
    for (const k of keys) {
        let resolveReady!: () => void;
        const ready = new Promise<void>((r) => {
            resolveReady = r;
        });
        controllerReady.push(ready);

        const branch = new ReadableStream<T>({
            start(controller) {
                controllers.set(k, controller);
                resolveReady();
            },
            cancel() {
                // Per-branch cancel is non-fatal; other branches keep flowing.
                // Drop the controller so writes to it become no-ops.
                controllers.delete(k);
            },
        });
        result.set(k, branch);
    }

    const startReader = async (): Promise<void> => {
        // Wait until every branch has registered its controller.
        await Promise.all(controllerReady);

        const reader = source.getReader();
        try {
            while (true) {
                const { value, done } = await reader.read();
                if (done) break;
                const k = keyFn(value as T);
                if (!knownKeys.has(k)) {
                    if (onUnknownKey === "error") {
                        const err = new StreamRouteUnknownKeyError({
                            key: String(k),
                        });
                        for (const c of controllers.values()) c.error(err);
                        controllers.clear();
                        try {
                            await reader.cancel(err);
                        } catch {
                            // ignore cancel failures
                        }
                        return;
                    }
                    // 'drop' — silently discard.
                    continue;
                }
                const c = controllers.get(k);
                if (c !== undefined) {
                    c.enqueue(value as T);
                }
            }
            for (const c of controllers.values()) c.close();
            controllers.clear();
        } catch (err) {
            for (const c of controllers.values()) c.error(err);
            controllers.clear();
        } finally {
            try {
                reader.releaseLock();
            } catch {
                // already released
            }
        }
    };

    // Kick the reader on first pull from any branch. We bind a `pull` by
    // wrapping each branch in a tee-style trigger: simplest correct approach
    // is to start eagerly on the next microtask once all branches are
    // constructed. Since branches are constructed synchronously above, this
    // microtask fires after the caller has the Map in hand.
    queueMicrotask(() => {
        void startReader();
    });

    return result;
}

/**
 * Routes each chunk of `source` to exactly one of N output branches, keyed by
 * the result of `keyFn`. Returns a record keyed by the declared `keys`. Each
 * chunk is delivered to a single branch (no duplication, single-pass).
 *
 * Per-key ordering is preserved. On unknown key, behaviour depends on
 * `options.onUnknownKey` (`'drop'` default, or `'error'`).
 */
export function route<T, K extends string>(
    source: ReadableStream<T>,
    keyFn: (chunk: T) => K,
    keys: readonly K[],
    options?: RouteOptions
): Record<K, ReadableStream<T>> {
    const m = buildRoutedMap<T, K>(source, keyFn, keys, options);
    const out = {} as Record<K, ReadableStream<T>>;
    for (const k of keys) {
        out[k] = m.get(k) as ReadableStream<T>;
    }
    return out;
}

/**
 * Like `route`, but returns a `Map<K, ReadableStream<T>>` — supporting key
 * types beyond `string` (numbers, symbols, objects).
 */
export function routeBy<T, K>(
    source: ReadableStream<T>,
    keyFn: (chunk: T) => K,
    keys: readonly K[],
    options?: RouteOptions
): Map<K, ReadableStream<T>> {
    return buildRoutedMap<T, K>(source, keyFn, keys, options);
}