import _sodium from "libsodium-wrappers";

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.slice().buffer);
}

function createFrameProcessor(key: Uint8Array) {
    // Выделяем nonce один раз — переиспользуем на каждый фрейм
    const nonce = new Uint8Array(sodium.crypto_stream_chacha20_NONCEBYTES); // 8 bytes
    const nonceView = new DataView(nonce.buffer, 0, nonce.byteLength);

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

        // Безопасно записываем nonce через отдельный DataView
        nonceView.setUint32(0, (timestamp >>> 0) & 0xffffffff, false);
        nonceView.setUint32(4, ((timestamp / 0x100000000) >>> 0) & 0xffffffff, false);

        const output = sodium.crypto_stream_chacha20_xor(input, nonce, key);
        return output.buffer as ArrayBuffer;
    };
}

function createTransform(processFrame: (data: ArrayBuffer, timestamp: number) => ArrayBuffer) {
    let frames = 0;
    let slowFrames = 0;
    let total = 0;
    let max = 0;

    return new TransformStream<any, any>({
        transform(frame, controller) {
            const started = performance.now();

            try {
                // Передаём timestamp фрейма как nonce
                const ts = typeof frame.timestamp === "number" ? frame.timestamp : 0;
                frame.data = processFrame(frame.data, ts);
            } catch (e) {
                console.error("[E2EE] frame error:", e);
                controller.enqueue(frame);
                return;
            }

            const elapsed = performance.now() - started;
            frames++;
            total += elapsed;
            if (elapsed > max) max = elapsed;
            if (elapsed > 2) slowFrames++;

            if (frames >= 200) {
                console.info("[E2EE stats]", {
                    avgMs: +(total / frames).toFixed(3),
                    maxMs: +max.toFixed(3),
                    slowFrames,
                });
                frames = 0;
                slowFrames = 0;
                total = 0;
                max = 0;
            }

            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);
    if (key.byteLength < sodium.crypto_stream_chacha20_KEYBYTES) {
        throw new Error(`Key must be at least ${sodium.crypto_stream_chacha20_KEYBYTES} 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, sodium.crypto_stream_chacha20_KEYBYTES));

    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);
    if (key.byteLength < sodium.crypto_stream_chacha20_KEYBYTES) {
        throw new Error(`Key must be at least ${sodium.crypto_stream_chacha20_KEYBYTES} 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, sodium.crypto_stream_chacha20_KEYBYTES));

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