import { sjclcodec } from "./codecBytes";
import sjcl, { SjclCipher } from "sjcl";

import { buf2hex } from "./utils";

interface EncryptionParams {
  ts: number;
  cs: number;
}

/**
 * AES-GCM TransformStream
 */
export class StreamEncryption {
  readonly #tagSize: number;
  readonly #chunkSize: number;
  readonly #aes: SjclCipher;
  readonly #key: sjcl.BitArray;
  readonly #iv: sjcl.BitArray;

  constructor(
    key: string | sjcl.BitArray | undefined,
    iv: string | sjcl.BitArray | undefined,
    params?: EncryptionParams | string,
  ) {
    if (!key && !iv) {
      key = buf2hex(globalThis.crypto.getRandomValues(new Uint8Array(16)));
      iv = buf2hex(globalThis.crypto.getRandomValues(new Uint8Array(12)));
    }
    if (typeof key === "string" && typeof iv === "string") {
      key = sjcl.codec.hex.toBits(key);
      iv = sjcl.codec.hex.toBits(iv);
    } else if (!Array.isArray(key) || !Array.isArray(iv)) {
      throw "Key and IV must be hex string or bitArray";
    }
    if (typeof params === "string") {
      params = JSON.parse(params) as EncryptionParams;
    }

    this.#tagSize = params?.ts ?? 128;
    this.#chunkSize = params?.cs ?? 1024 * 1024 * 10;
    this.#aes = new sjcl.cipher.aes(key);
    this.#key = key;
    this.#iv = iv;

    console.log(
      `ts=${this.#tagSize}, cs=${this.#chunkSize}, key=${key}, iv=${this.#iv}`,
    );
  }

  /**
   * Return formatted encryption key
   */
  getKey() {
    return `${sjcl.codec.hex.fromBits(this.#key)}:${sjcl.codec.hex.fromBits(
      this.#iv,
    )}`;
  }

  /**
   * Get encryption params
   */
  getParams() {
    return {
      ts: this.#tagSize,
      cs: this.#chunkSize,
    };
  }

  /**
   * Get encryption TransformStream
   */
  getEncryptionTransform() {
    return this._getCryptoStream(0);
  }

  /**
   * Get decryption TransformStream
   */
  getDecryptionTransform() {
    return this._getCryptoStream(1);
  }

  _getCryptoStream(mode: number) {
    let offset = 0;
    let buffer = new Uint8Array(
      this.#chunkSize + (mode === 1 ? this.#tagSize / 8 : 0),
    );
    return new TransformStream(
      {
        transform: async (chunk, controller) => {
          chunk = await chunk;
          try {
            let toBuffer = Math.min(
              chunk.byteLength,
              buffer.byteLength - offset,
            );
            buffer.set(chunk.slice(0, toBuffer), offset);
            offset += toBuffer;

            if (offset === buffer.byteLength) {
              let buff = sjclcodec.toBits(buffer);
              let encryptedBuf = sjclcodec.fromBits(
                mode === 0
                  ? sjcl.mode.gcm.encrypt(
                      this.#aes,
                      buff,
                      this.#iv,
                      [],
                      this.#tagSize,
                    )
                  : sjcl.mode.gcm.decrypt(
                      this.#aes,
                      buff,
                      this.#iv,
                      [],
                      this.#tagSize,
                    ),
              );
              controller.enqueue(new Uint8Array(encryptedBuf));

              offset = chunk.byteLength - toBuffer;
              buffer.set(chunk.slice(toBuffer));
            }
          } catch (e) {
            console.error(e);
            throw e;
          }
        },
        flush: (controller) => {
          let lastBuffer = buffer.slice(0, offset);
          let buff = sjclcodec.toBits(lastBuffer);
          let encryptedBuf = sjclcodec.fromBits(
            mode === 0
              ? sjcl.mode.gcm.encrypt(
                  this.#aes,
                  buff,
                  this.#iv,
                  [],
                  this.#tagSize,
                )
              : sjcl.mode.gcm.decrypt(
                  this.#aes,
                  buff,
                  this.#iv,
                  [],
                  this.#tagSize,
                ),
          );
          controller.enqueue(new Uint8Array(encryptedBuf));
        },
      },
      {
        highWaterMark: this.#chunkSize,
      },
    );
  }
}