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);
}

/**
 * Переиспользуемый процессор фреймов.
 * Один экземпляр на sender/receiver — нет аллокаций на каждый фрейм.
 */
function createFrameProcessor(key: Uint8Array) {
    // Переиспользуемые буферы — не создаём новые на каждый фрейм
    const nonce = new Uint8Array(12);
    const nonceView = new DataView(nonce.buffer);

    return function processFrame(data: ArrayBuffer): ArrayBuffer {
        // Переиспользуем nonce буфер
        nonce.fill(0);
        nonceView.setUint32(8, (data.byteLength ^ (data.byteLength << 8)) >>> 0, false);

        const input = new Uint8Array(data);
        const output = chacha20(key, nonce, input);
        return output.buffer as ArrayBuffer;
    };
}

function createTransform(processFrame: (data: ArrayBuffer) => ArrayBuffer) {
    return new TransformStream<any, any>({
        transform(frame, controller) {
            try {
                const start = performance.now();
                frame.data = processFrame(frame.data);
                const elapsed = performance.now() - start;
                if (elapsed > 1) {
                    console.warn(`E2EE slow frame: ${elapsed.toFixed(2)}ms`);
                }
                controller.enqueue(frame);
            } catch (e) {
                // не рвём поток — пропускаем фрейм как есть
                console.error("E2EE frame failed:", e);
                controller.enqueue(frame);
            }
        }
    });
}

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 not available on RTCRtpSender");

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

    readable
        .pipeThrough(createTransform(processFrame))
        .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 not available on RTCRtpReceiver");

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

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