import { resolver } from '@rocicorp/resolver';
import { pipeline, Readable, Transform, Writable, } from 'node:stream';
import { createWebSocketStream, } from 'ws';
import { BigIntJSON } from "../../../shared/src/bigint-json.js";
import { Queue } from "../../../shared/src/queue.js";
import * as v from "../../../shared/src/valita.js";
import { Subscription } from "./subscription.js";
import { closeWithError, expectPingsForLiveness, sendPingsForLiveness, } from "./ws.js";
// Consistent with Postgres keepalives, and shorter than the
// commonly used default idle timeout of 1 minute.
const PING_INTERVAL_MS = 30_000;
/**
 * Back-pressure-aware transformation of a WebSocket into
 * upstream and downstream {@link Subscription} objects.
 */
// TODO: Change {@link streamIn} and {@link streamOut} to use this
//       under the covers so that internal communication is also
//       responsive to backpressure.
export function stream(lc, ws, inSchema, outOptions = {}, inOptions = {}, streamOptions = {}) {
    const endpoint = ws.url ?? 'client';
    function close(err) {
        if (ws.readyState !== ws.CLOSED && ws.readyState !== ws.CLOSING) {
            if (err) {
                closeWithError(lc, ws, err);
            }
            else {
                lc.info?.(`closing connection to ${endpoint}`);
                ws.close();
            }
        }
    }
    const instream = Subscription.create({
        ...inOptions,
        cleanup: (unconsumed, err) => {
            inOptions.cleanup?.(unconsumed, err);
            close(err);
        },
    });
    const outstream = Subscription.create({
        ...outOptions,
        cleanup: (unconsumed, err) => {
            outOptions.cleanup?.(unconsumed, err);
            close(err);
        },
    });
    const duplex = createWebSocketStream(ws, {
        ...streamOptions,
        decodeStrings: false,
    });
    // Outgoing transform.
    function streamOut() {
        // Mainly used for verifying that back-pressure kicks in tests.
        duplex.on('drain', () => lc.debug?.(`drained messages to ${endpoint}`));
        pipeline(Readable.from(outstream), new Transform({
            objectMode: true,
            transform: (msg, _encoding, callback) => callback(null, BigIntJSON.stringify(msg)),
        }), duplex, err => (err ? outstream.fail(err) : outstream.cancel()));
    }
    if (ws.readyState === ws.CONNECTING) {
        ws.on('open', () => {
            lc.info?.(`connected to ${endpoint}`);
            streamOut();
        });
    }
    else {
        streamOut();
    }
    // Incoming transform.
    pipe(duplex, instream, chunk => {
        const json = BigIntJSON.parse(chunk.toString());
        return v.parse(json, inSchema, 'passthrough');
    });
    sendPingsForLiveness(lc, ws, PING_INTERVAL_MS);
    return { outstream, instream };
}
export function pipe(source, sink, parse) {
    pipeline(source, new Writable({
        decodeStrings: false,
        write: (chunk, _encoding, callback) => {
            let msg;
            try {
                if ((msg = parse(chunk)) === null) {
                    callback();
                    return;
                }
            }
            catch (err) {
                callback(ensureError(err));
                return;
            }
            const { result } = sink.push(msg);
            // Inbound backpressure is exerted by unconsumed
            // messages in the subscription.
            void result.then(() => callback(), err => callback(ensureError(err)));
        },
        destroy: (err, callback) => {
            if (err) {
                sink.fail(ensureError(err));
            }
            // Otherwise, final will handle the cancel.
            callback();
        },
        final: callback => {
            sink.cancel();
            callback();
        },
    }), err => (err ? sink.fail(err) : sink.cancel()));
}
function ensureError(err) {
    return err instanceof Error ? err : new Error(String(err));
}
const ackSchema = v.object({ ack: v.number() });
export async function streamOut(lc, source, sink) {
    sendPingsForLiveness(lc, sink, PING_INTERVAL_MS);
    const closer = WebSocketCloser.forSource(lc, sink, source);
    const acks = new Queue();
    sink.addEventListener('message', ({ data }) => {
        try {
            if (typeof data !== 'string') {
                throw new Error('Expected string message');
            }
            acks.enqueue(v.parse(JSON.parse(data), ackSchema));
        }
        catch (e) {
            lc.error?.(`error parsing ack`, e);
            closer.close(e);
        }
    });
    try {
        let nextID = 0;
        const { pipeline } = source;
        if (pipeline) {
            lc.debug?.(`started pipelined outbound stream`);
            for await (const { value: msg, consumed } of pipeline) {
                const id = ++nextID;
                const data = BigIntJSON.stringify({ msg, id });
                // Enable for debugging. Otherwise too verbose.
                // lc.debug?.(`pipelining`, data);
                sink.send(data);
                void (async () => {
                    const { ack } = await acks.dequeue();
                    // lc.debug?.(`received ack`, ack);
                    if (ack !== id) {
                        throw new Error(`Unexpected ack for ${id}: ${ack}`);
                    }
                    consumed();
                })();
            }
        }
        else {
            lc.debug?.(`started synchronous outbound stream`);
            for await (const msg of source) {
                const id = ++nextID;
                const data = BigIntJSON.stringify({ msg, id });
                // Enable for debugging. Otherwise too verbose.
                // lc.debug?.(`sending`, data);
                sink.send(data);
                const { ack } = await acks.dequeue();
                if (ack !== id) {
                    throw new Error(`Unexpected ack for ${id}: ${ack}`);
                }
            }
        }
        closer.close();
    }
    catch (e) {
        closer.close(e);
    }
}
export async function streamIn(lc, source, schema) {
    expectPingsForLiveness(lc, source, PING_INTERVAL_MS);
    const streamedSchema = v.object({
        msg: schema,
        id: v.number(),
    });
    const sink = new Subscription({
        consumed: ({ id }) => source.send(JSON.stringify({ ack: id })),
        cleanup: () => closer.close(),
    }, ({ msg }) => msg);
    const closer = WebSocketCloser.forSink(lc, source, sink, handleMessage);
    function handleMessage(event) {
        const data = event.data.toString();
        if (!sink.active) {
            lc.warn?.('dropping ws message received after close', data);
            return;
        }
        try {
            const value = BigIntJSON.parse(data);
            const msg = v.parse(value, streamedSchema, 'passthrough');
            // Enable for debugging. Otherwise too verbose.
            // lc.debug?.(`received`, data);
            sink.push(msg);
        }
        catch (e) {
            closer.close(e);
        }
    }
    await closer.connected;
    return sink;
}
class WebSocketCloser {
    #lc;
    #ws;
    #closeStream;
    #messageHandler;
    #connected = resolver();
    get connected() {
        return this.#connected.promise;
    }
    static forSource(lc, ws, stream) {
        // If the websocket is closed, call cancel() to notify the Source of
        // any unconsumed messages.
        return new WebSocketCloser(lc, ws, () => stream.cancel());
    }
    static forSink(lc, ws, stream, messageHandler) {
        // If the websocket is closed, call end() to allow the downstream Sink
        // to process any pending messages before closing the stream.
        return new WebSocketCloser(lc, ws, () => stream.end(), messageHandler);
    }
    constructor(lc, ws, closeStream, messageHandler) {
        this.#lc = lc;
        this.#ws = ws;
        this.#closeStream = closeStream;
        this.#messageHandler = messageHandler ?? null;
        ws.addEventListener('open', this.#handleOpen);
        ws.addEventListener('close', this.#handleClose);
        ws.addEventListener('error', this.#handleError);
        if (this.#messageHandler) {
            ws.addEventListener('message', this.#messageHandler);
        }
        switch (ws.readyState) {
            case ws.CONNECTING:
                break; // expected for new connections. resolve or reject in handlers.
            case ws.OPEN:
                this.#connected.resolve();
                break;
            default:
                this.#connected.reject(new Error(`websocket already in state ${ws.readyState}`));
                break;
        }
    }
    get #conn() {
        return 'connection' + (this.#ws.url ? ` to ${this.#ws.url}` : '');
    }
    #handleOpen = () => {
        this.#lc.info?.(`${this.#conn} established`);
        this.#connected.resolve();
    };
    #handleClose = (e) => {
        const { code, reason, wasClean } = e;
        this.#lc.info?.(`${this.#conn} closed`, {
            code,
            reason,
            wasClean,
        });
        this.close();
        this.#connected.reject(`${this.#conn} closed with code ${code}`);
    };
    #handleError = ({ message, error }) => {
        if (this.#ws.readyState === this.#ws.OPEN) {
            this.#lc.error?.(`error in ${this.#conn}`, message, error);
        }
        this.#connected.reject(error);
    };
    close(err) {
        if (err) {
            this.#lc.error?.(`closing stream with error`, err);
        }
        this.#closeStream();
        if (!this.closed()) {
            this.#ws.close();
        }
    }
    closed() {
        return (this.#ws.readyState === this.#ws.CLOSED ||
            this.#ws.readyState === this.#ws.CLOSING);
    }
}
//# sourceMappingURL=streams.js.map