function toArrayBuffer(src: Buffer | Uint8Array): ArrayBuffer {
  const u8 = src instanceof Uint8Array ? src : new Uint8Array(src);
  return u8.buffer.slice(u8.byteOffset, u8.byteOffset + u8.byteLength) as ArrayBuffer;
}

type KeyInput = Buffer | Uint8Array;


async function importAesCtrKey(input: KeyInput): Promise<CryptoKey> {
console.info("Importing AES-CTR key for E2EE:", Buffer.from(input).toString('hex'));
  const keyBytes = toArrayBuffer(input);
  if (keyBytes.byteLength !== 32) {
    throw new Error(`E2EE key must be 32 bytes, got ${keyBytes.byteLength}`);
  }

  return crypto.subtle.importKey(
    "raw",
    keyBytes,
    { name: "AES-CTR" },
    false,
    ["encrypt", "decrypt"]
  );
}

function toBigIntTs(ts: unknown): bigint {
  if (typeof ts === "bigint") return ts;
  if (typeof ts === "number") return BigInt(ts);
  return 0n;
}

/**
 * 16-byte counter:
 * [0..3]   direction marker
 * [4..11]  frame timestamp
 * [12..15] reserved
 */
function buildCounter(direction: number, timestamp: unknown): ArrayBuffer {
  const iv = new Uint8Array(16);
  const dv = new DataView(iv.buffer);
  dv.setUint32(0, direction >>> 0, false);
  dv.setBigUint64(4, toBigIntTs(timestamp), false);
  dv.setUint32(12, 0, false);
  return toArrayBuffer(iv);
}

export async function attachSenderE2EE(sender: RTCRtpSender, keyInput: KeyInput): Promise<void> {
  const key = await importAesCtrKey(keyInput);

  const anySender = sender as unknown as {
    createEncodedStreams?: () => { readable: ReadableStream<any>; writable: WritableStream<any> };
  };

  if (!anySender.createEncodedStreams) {
    throw new Error("createEncodedStreams is not available on RTCRtpSender");
  }

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

  const enc = new TransformStream<any, any>({
    async transform(frame, controller) {
      const counter = buildCounter(1, frame.timestamp);
      const encrypted = await crypto.subtle.encrypt(
        { name: "AES-CTR", counter, length: 64 },
        key,
        frame.data
      );
      frame.data = encrypted; // same length
      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> {
  const key = await importAesCtrKey(keyInput);

  const anyReceiver = receiver as unknown as {
    createEncodedStreams?: () => { readable: ReadableStream<any>; writable: WritableStream<any> };
  };

  if (!anyReceiver.createEncodedStreams) {
    throw new Error("createEncodedStreams is not available on RTCRtpReceiver");
  }

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

  const dec = new TransformStream<any, any>({
    async transform(frame, controller) {
      const counter = buildCounter(1, frame.timestamp);
      const decrypted = await crypto.subtle.decrypt(
        { name: "AES-CTR", counter, length: 64 },
        key,
        frame.data
      );
      frame.data = decrypted; // same length
      controller.enqueue(frame);
    }
  });

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