/** copied from https://github.com/langchain-ai/langgraphjs/tree/main/libs/sdk/src/utils */
const CR = "\r".charCodeAt(0);
const LF = "\n".charCodeAt(0);
const NULL = "\0".charCodeAt(0);
const COLON = ":".charCodeAt(0);
const SPACE = " ".charCodeAt(0);

const TRAILING_NEWLINE = [CR, LF];

export function BytesLineDecoder() {
    let buffer: Uint8Array[] = [];
    let trailingCr = false;

    return new TransformStream<Uint8Array, Uint8Array>({
        start() {
            buffer = [];
            trailingCr = false;
        },

        transform(chunk, controller) {
            // See https://docs.python.org/3/glossary.html#term-universal-newlines
            let text = chunk;

            // Handle trailing CR from previous chunk
            if (trailingCr) {
                text = joinArrays([[CR], text]);
                trailingCr = false;
            }

            // Check for trailing CR in current chunk
            if (text.length > 0 && text.at(-1) === CR) {
                trailingCr = true;
                text = text.subarray(0, -1);
            }

            if (!text.length) return;
            const trailingNewline = TRAILING_NEWLINE.includes(text.at(-1)!);

            const lastIdx = text.length - 1;
            const { lines } = text.reduce<{ lines: Uint8Array[]; from: number }>(
                (acc, cur, idx) => {
                    if (acc.from > idx) return acc;

                    if (cur === CR || cur === LF) {
                        acc.lines.push(text.subarray(acc.from, idx));
                        if (cur === CR && text[idx + 1] === LF) {
                            acc.from = idx + 2;
                        } else {
                            acc.from = idx + 1;
                        }
                    }

                    if (idx === lastIdx && acc.from <= lastIdx) {
                        acc.lines.push(text.subarray(acc.from));
                    }

                    return acc;
                },
                { lines: [], from: 0 }
            );

            if (lines.length === 1 && !trailingNewline) {
                buffer.push(lines[0]);
                return;
            }

            if (buffer.length) {
                // Include existing buffer in first line
                buffer.push(lines[0]);
                lines[0] = joinArrays(buffer);
                buffer = [];
            }

            if (!trailingNewline) {
                // If the last segment is not newline terminated,
                // buffer it for the next chunk
                if (lines.length) buffer = [lines.pop()!];
            }

            // Enqueue complete lines
            for (const line of lines) {
                controller.enqueue(line);
            }
        },

        flush(controller) {
            if (buffer.length) {
                controller.enqueue(joinArrays(buffer));
            }
        },
    });
}

interface StreamPart {
    id: string | undefined;
    event: string;
    data: unknown;
}

export function SSEDecoder() {
    let event = "";
    let data: Uint8Array[] = [];
    let lastEventId = "";
    let retry: number | null = null;

    const decoder = new TextDecoder();

    return new TransformStream<Uint8Array, StreamPart>({
        transform(chunk, controller) {
            // Handle empty line case
            if (!chunk.length) {
                if (!event && !data.length && !lastEventId && retry == null) return;

                const sse = {
                    id: lastEventId || undefined,
                    event,
                    data: data.length ? decodeArraysToJson(decoder, data) : null,
                };

                // NOTE: as per the SSE spec, do not reset lastEventId
                event = "";
                data = [];
                retry = null;

                controller.enqueue(sse);
                return;
            }

            // Ignore comments
            if (chunk[0] === COLON) return;

            const sepIdx = chunk.indexOf(COLON);
            if (sepIdx === -1) return;

            const fieldName = decoder.decode(chunk.subarray(0, sepIdx));
            let value = chunk.subarray(sepIdx + 1);
            if (value[0] === SPACE) value = value.subarray(1);

            if (fieldName === "event") {
                event = decoder.decode(value);
            } else if (fieldName === "data") {
                data.push(value);
            } else if (fieldName === "id") {
                if (value.indexOf(NULL) === -1) lastEventId = decoder.decode(value);
            } else if (fieldName === "retry") {
                const retryNum = Number.parseInt(decoder.decode(value), 10);
                if (!Number.isNaN(retryNum)) retry = retryNum;
            }
        },

        flush(controller) {
            if (event) {
                controller.enqueue({
                    id: lastEventId || undefined,
                    event,
                    data: data.length ? decodeArraysToJson(decoder, data) : null,
                });
            }
        },
    });
}

function joinArrays(data: ArrayLike<number>[]) {
    const totalLength = data.reduce((acc, curr) => acc + curr.length, 0);
    const merged = new Uint8Array(totalLength);
    let offset = 0;
    for (const c of data) {
        merged.set(c, offset);
        offset += c.length;
    }
    return merged;
}

function decodeArraysToJson(decoder: TextDecoder, data: ArrayLike<number>[]) {
    return JSON.parse(decoder.decode(joinArrays(data)));
}