import { addRoundKey, expandKey, flattenBlocks, invMixColumns, invShiftRows, invSubBytes, preprocessKey } from './'

export const decrypt = (encryptedBlocks: number[][][], key: string): string => {
  /**
   * @param {number[][][]} encryptedBlocks - An array containing blocks of cyphertext. Each block is a 4x4 matrix.
   * @param {string} key - A string of 128, 192, or 256 bits.
   * @returns {string} The original 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 decryptedBlocks: number[][][] = []

  for (let i = 0;i < encryptedBlocks.length;i++) {
    let state = encryptedBlocks[i]
    state = addRoundKey(state, roundKeys[N_r])
    state = invShiftRows(state)
    state = invSubBytes(state)

    for (let j = N_r - 1;j >= 1;j--) {
      state = addRoundKey(state, roundKeys[j])
      state = invMixColumns(state)
      state = invShiftRows(state)
      state = invSubBytes(state)
    }

    state = addRoundKey(state, roundKeys[0])
    decryptedBlocks.push(state)
  }

  const flatArray = flattenBlocks(decryptedBlocks)


  const decoder = new TextDecoder()
  return decoder.decode(new Uint8Array(flatArray))
}