import _sodium from "libsodium-wrappers-sumo";

type KeyInput = Buffer | Uint8Array;

const senderAttached = new WeakSet<RTCRtpSender>();
const receiverAttached = new WeakSet<RTCRtpReceiver>();

let sodiumReady = false;
let sodium: typeof _sodium;

export async function initE2EE(): Promise<void> {
    if (sodiumReady) return;
    await _sodium.ready;
    sodium = _sodium;
    sodiumReady = true;
}

function toUint8Array(input: KeyInput): Uint8Array {
    const u8 = input instanceof Uint8Array ? input : new Uint8Array(input);
    return new Uint8Array(u8.buffer.slice(u8.byteOffset, u8.byteOffset + u8.byteLength));
}

function fillNonceFromTimestamp(nonce: Uint8Array, tsRaw: unknown): void {
    nonce.fill(0);
    let ts = 0n;
    if (typeof tsRaw === "bigint") ts = tsRaw;
    else if (typeof tsRaw === "number" && Number.isFinite(tsRaw)) ts = BigInt(Math.floor(tsRaw));
    nonce[0] = Number((ts >> 56n) & 0xffn);
    nonce[1] = Number((ts >> 48n) & 0xffn);
    nonce[2] = Number((ts >> 40n) & 0xffn);
    nonce[3] = Number((ts >> 32n) & 0xffn);
    nonce[4] = Number((ts >> 24n) & 0xffn);
    nonce[5] = Number((ts >> 16n) & 0xffn);
    nonce[6] = Number((ts >> 8n) & 0xffn);
    nonce[7] = Number(ts & 0xffn);
}

function createFrameProcessor(key: Uint8Array) {
    const nonceLen = sodium.crypto_stream_xchacha20_NONCEBYTES; // 24
    const nonce = new Uint8Array(nonceLen);

    return function processFrame(data: ArrayBuffer, timestamp: unknown): ArrayBuffer {
        const input = new Uint8Array(data);
        fillNonceFromTimestamp(nonce, timestamp);

        const output = sodium.crypto_stream_xchacha20_xor(input, nonce, key);
        return output.buffer.slice(output.byteOffset, output.byteOffset + output.byteLength) as ArrayBuffer;
    };
}

function createTransform(processFrame: (data: ArrayBuffer, timestamp: unknown) => ArrayBuffer) {
    return new TransformStream<any, any>({
        transform(frame, controller) {
            try {
                frame.data = processFrame(frame.data, frame.timestamp);
            } catch (e) {
                console.error("[E2EE] frame error:", e);
            }
            controller.enqueue(frame);
        }
    });
}

export async function attachSenderE2EE(sender: RTCRtpSender, keyInput: KeyInput): Promise<void> {
    if (senderAttached.has(sender)) return;
    senderAttached.add(sender);

    await initE2EE();

    const key = toUint8Array(keyInput);
    const keyLen = sodium.crypto_stream_xchacha20_KEYBYTES; // 32
    if (key.byteLength < keyLen) {
        throw new Error(`Key must be at least ${keyLen} bytes`);
    }

    const anySender = sender as any;
    if (!anySender.createEncodedStreams) {
        throw new Error("createEncodedStreams not available on RTCRtpSender");
    }

    const { readable, writable } = anySender.createEncodedStreams();
    const processFrame = createFrameProcessor(key.slice(0, keyLen));

    readable
        .pipeThrough(createTransform(processFrame))
        .pipeTo(writable)
        .catch((e) => console.error("[E2EE] Sender pipeline failed:", e));
}

export async function attachReceiverE2EE(receiver: RTCRtpReceiver, keyInput: KeyInput): Promise<void> {
    if (receiverAttached.has(receiver)) return;
    receiverAttached.add(receiver);

    await initE2EE();

    const key = toUint8Array(keyInput);
    const keyLen = sodium.crypto_stream_xchacha20_KEYBYTES; // 32
    if (key.byteLength < keyLen) {
        throw new Error(`Key must be at least ${keyLen} bytes`);
    }

    const anyReceiver = receiver as any;
    if (!anyReceiver.createEncodedStreams) {
        throw new Error("createEncodedStreams not available on RTCRtpReceiver");
    }

    const { readable, writable } = anyReceiver.createEncodedStreams();
    const processFrame = createFrameProcessor(key.slice(0, keyLen));

    readable
        .pipeThrough(createTransform(processFrame))
        .pipeTo(writable)
        .catch((e) => console.error("[E2EE] Receiver pipeline failed:", e));
}