import { preprocessPlaintext, preprocessKey, expandKey, addRoundKey, subBytes } from './'
import { mixColumns } from './mixColumns'
import { shiftRows } from './shiftRows'

export const encrypt = (plaintext: string, key: string): number[][][] => {
  /**
   * @param {string} plaintext - The plaintext to encrypt.
   * @param {string} key - A string of 128, 192, or 256 bits.
   * @returns {string} The encrypted cyphertext.
   */

  const plaintextBlocks = preprocessPlaintext(plaintext)
  const keyBlock = preprocessKey(key)
  const roundKeys = expandKey(keyBlock)

  const keySize = key.length // Key size in bytes (16 for 128-bit, 24 for 192-bit, 32 for 256-bit)
  const N_k = keySize / 4 // Number of 32-bit words in the key (4 for 128-bit, 6 for 192-bit, 8 for 256-bit)
  const N_r = N_k + 6

  const encryptedBlocks: number[][][] = []

  for (let i = 0;i < plaintextBlocks.length;i++) {
    let state = plaintextBlocks[i]
    state = addRoundKey(state, roundKeys[0])
    for (let j = 1;j <= N_r - 1;j++) {
      state = subBytes(state)
      state = shiftRows(state)
      state = mixColumns(state)
      state = addRoundKey(state, roundKeys[j])
    }
    state = subBytes(state)
    state = shiftRows(state)
    state = addRoundKey(state, roundKeys[N_r])
    encryptedBlocks.push(state)
  }
  
  return encryptedBlocks
}