import { chacha20 } from "@noble/ciphers/chacha";

type KeyInput = Buffer | Uint8Array;

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

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

function buildNonce(timestamp: unknown): Uint8Array {
    const nonce = new Uint8Array(12);
    const ts = typeof timestamp === "number"
        ? timestamp
        : typeof timestamp === "bigint"
            ? Number(timestamp)
            : 0;
    new DataView(nonce.buffer).setUint32(8, ts >>> 0, false);
    return nonce;
}

function processFrame(data: ArrayBuffer, key: Uint8Array, timestamp: unknown): ArrayBuffer {
    const nonce = buildNonce(timestamp);
    const input = new Uint8Array(data);
    // ChaCha20 симметричный: encrypt === decrypt, тот же размер
    const output = chacha20(key, nonce, input);
    return output.buffer as ArrayBuffer;
}

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

    const key = toUint8Array(keyInput);
    if (key.byteLength !== 32) {
        throw new Error(`E2EE key must be 32 bytes, got ${key.byteLength}`);
    }

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

    const { readable, writable } = anySender.createEncodedStreams();

    const enc = new TransformStream<any, any>({
        // Синхронный transform — нет async, нет накопления очереди
        transform(frame, controller) {
            try {
                frame.data = processFrame(frame.data, key, frame.timestamp);
                controller.enqueue(frame);
            } catch (e) {
                console.error("Sender E2EE frame failed:", e);
                controller.enqueue(frame);
            }
        }
    });

    readable.pipeThrough(enc).pipeTo(writable).catch((e) => {
        console.error("Sender E2EE pipeline failed:", e);
    });
}

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

    const key = toUint8Array(keyInput);
    if (key.byteLength !== 32) {
        throw new Error(`E2EE key must be 32 bytes, got ${key.byteLength}`);
    }

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

    const { readable, writable } = anyReceiver.createEncodedStreams();

    const dec = new TransformStream<any, any>({
        // Синхронный transform — нет async, нет накопления очереди
        transform(frame, controller) {
            try {
                frame.data = processFrame(frame.data, key, frame.timestamp);
                controller.enqueue(frame);
            } catch (e) {
                console.error("Receiver E2EE frame failed:", e);
                controller.enqueue(frame);
            }
        }
    });

    readable.pipeThrough(dec).pipeTo(writable).catch((e) => {
        console.error("Receiver E2EE pipeline failed:", e);
    });
}