import { promiseWithResolvers } from "../../../utils/promiseWithResolvers.js";
export const createMergeStream = () => {
    const list = [];
    let sealed = false;
    let controller;
    let currentPull;
    const handlePull = (item) => {
        if (!item.promise) {
            // TODO for most streams, we can directly pipeTo to avoid the microTask queue
            // add an option to eagerly pipe the stream to the merge stream
            // ideally, using assistant-stream w sync run method + piping to a sync WritableStream runs in the same microtask
            // this is useful because we often use AssistantStreams internally as a serialization utility, e. g. AssistantTransformStream
            // idea: avoid reader.read() by instead using a WritableStream & if (!hasPendingPull) await waitForPull()?
            item.promise = item.reader
                .read()
                .then(({ done, value }) => {
                item.promise = undefined;
                if (done) {
                    list.splice(list.indexOf(item), 1);
                    if (sealed && list.length === 0) {
                        controller.close();
                    }
                }
                else {
                    controller.enqueue(value);
                }
                currentPull?.resolve();
                currentPull = undefined;
            })
                .catch((e) => {
                console.error(e);
                list.forEach((item) => {
                    item.reader.cancel();
                });
                list.length = 0;
                controller.error(e);
                currentPull?.reject(e);
                currentPull = undefined;
            });
        }
    };
    const readable = new ReadableStream({
        start(c) {
            controller = c;
        },
        pull() {
            currentPull = promiseWithResolvers();
            list.forEach((item) => {
                handlePull(item);
            });
            return currentPull.promise;
        },
        cancel() {
            list.forEach((item) => {
                item.reader.cancel();
            });
            list.length = 0;
        },
    });
    return {
        readable,
        isSealed() {
            return sealed;
        },
        seal() {
            sealed = true;
            if (list.length === 0)
                controller.close();
        },
        addStream(stream) {
            if (sealed)
                throw new Error("Cannot add streams after the run callback has settled.");
            const item = { reader: stream.getReader() };
            list.push(item);
            handlePull(item);
        },
        enqueue(chunk) {
            this.addStream(new ReadableStream({
                start(c) {
                    c.enqueue(chunk);
                    c.close();
                },
            }));
        },
    };
};
// TODO
// export class SpanContainerMerger {
//   public get isSealed() {
//     return this.mergeStream.isSealed();
//   }
//   public get readable() {
//     return this.mergeStream.readable;
//   }
//   private subAllocator = new Counter();
//   private mergeStream = createMergeStream();
//   constructor() {
//     // id 0 is auto allocated
//     this.subAllocator.up();
//   }
//   add(stream: ReadableStream<AssistantStreamChunk>) {
//     this.mergeStream.addStream(
//       stream.pipeThrough(new SpanParentEncoder(this.subAllocator)),
//     );
//   }
//   enqueue(chunk: AssistantStreamChunk & { parentId: 0 }) {
//     this.mergeStream.addStream(
//       new ReadableStream({
//         start(c) {
//           c.enqueue(chunk);
//           c.close();
//         },
//       }),
//     );
//   }
//   seal() {
//     this.mergeStream.seal();
//   }
// }
// export class SpanContainerSplitter {
//   public writable;
//   private isSealed = false;
//   private writers = new Map<
//     number,
//     WritableStreamDefaultWriter<AssistantStreamChunk>
//   >();
//   private closeTasks: Promise<void>[] = [];
//   private allocator = new Counter();
//   private subAllocator = new Counter();
//   constructor() {
//     // id 0 is auto-allocated
//     this.allocator.up();
//     this.writable = new WritableStream({
//       write: (chunk) => {
//         const { type, parentId } = chunk;
//         const writer = this.writers.get(parentId);
//         if (writer === undefined) throw new Error("Parent id not found");
//         writer.write(chunk);
//         if (type === "span") {
//           // allocate a new span id
//           this.writers.set(this.allocator.up(), writer);
//         }
//         if (type === "finish") {
//           this.writers.delete(parentId);
//           writer.close();
//           if (this.writers.size === 0) {
//             const closeTask = this.writable.close();
//             this.closeTasks.push(closeTask);
//             closeTask.then(() => {
//               this.closeTasks.splice(this.closeTasks.indexOf(closeTask), 1);
//             });
//           }
//         }
//       },
//       close: async () => {
//         if (this.writers.size > 0) throw new Error("Not all writers closed");
//         // await and throw on any errors
//         await Promise.all(this.closeTasks);
//       },
//     });
//   }
//   add(stream: WritableStream<AssistantStreamChunk>) {
//     if (this.isSealed) throw new Error("Cannot add streams after sealing");
//     const decoder = new SpanParentDecoder(this.subAllocator);
//     decoder.readable.pipeTo(stream);
//     this.writers.set(this.allocator.up(), decoder.writable.getWriter());
//   }
//   seal() {
//     this.isSealed = true;
//     if (this.writers.size === 0) this.writable.close();
//   }
// }
//# sourceMappingURL=merge.js.map