import { Buffer } from "buffer";
import {IHandshake} from "./@types/handshake-interface";
import {NOISE_MSG_MAX_LENGTH_BYTES, NOISE_MSG_MAX_LENGTH_BYTES_WITHOUT_TAG} from "./constants";

interface IReturnEncryptionWrapper {
  (source: Iterable<Uint8Array>): AsyncIterableIterator<Uint8Array>;
}

// Returns generator that encrypts payload from the user
export function encryptStream(handshake: IHandshake): IReturnEncryptionWrapper {
  return async function * (source) {
    for await (const chunk of source) {
      const chunkBuffer = Buffer.from(chunk.buffer, chunk.byteOffset, chunk.length);

      for (let i = 0; i < chunkBuffer.length; i += NOISE_MSG_MAX_LENGTH_BYTES_WITHOUT_TAG) {
        let end = i + NOISE_MSG_MAX_LENGTH_BYTES_WITHOUT_TAG;
        if (end > chunkBuffer.length) {
          end = chunkBuffer.length;
        }

        const data = handshake.encrypt(chunkBuffer.slice(i, end), handshake.session);
        yield data;
      }
    }
  }
}


// Decrypt received payload to the user
export function decryptStream(handshake: IHandshake): IReturnEncryptionWrapper {
  return async function * (source) {
    for await (const chunk of source) {
      const chunkBuffer = Buffer.from(chunk.buffer, chunk.byteOffset, chunk.length);

      for (let i = 0; i < chunkBuffer.length; i += NOISE_MSG_MAX_LENGTH_BYTES) {
        let end = i + NOISE_MSG_MAX_LENGTH_BYTES;
        if (end > chunkBuffer.length) {
          end = chunkBuffer.length;
        }

        const chunk = chunkBuffer.slice(i, end);
        const decrypted = await handshake.decrypt(chunk, handshake.session);
        yield decrypted;
      }
    }
  }
}