import { RCons, SBox } from '../constants'

const subWord = (word: number[]): number[] => {
  /**
   * @param {number[]} word - An an array containing 4 numbers, each number is a hex value represeting one byte.
   * @returns {number[]} An array containing the corresponding values in the S-box.
   */
  return word.map(byte => SBox[byte])
}

const rotWord = (word: number[]): number[] => {
  return word.slice(1).concat(word[0])
}

export const expandKey = (key: number[][]): number[][][] => {
  /**
   * @param {string} key - A string of 128, 192, or 256 bits.
   * @returns {string[]} An array of round keys.
   */
  const N_k = key[0].length // Number of columns in the key matrix (4 for 128-bit, 6 for 192-bit, 8 for 256-bit)
  const N_r = N_k + 6 // Number of rounds (10 for 128-bit, 12 for 192-bit, 14 for 256-bit)
  const totalKeys = 4 * (N_r + 1) // Total number of 4-byte words required

  const roundKeys: number[][] = []

  // Step 1: Copy the initial key as the first round key (N_k columns)
  for (let i = 0;i < N_k;i++) {
    roundKeys.push([
      key[0][i],
      key[1][i],
      key[2][i],
      key[3][i]
    ])
  }

  // Step 2: Expand the key to generate the full round keys
  for (let i = N_k;i < totalKeys;i++) {
    let temp = roundKeys[i - 1]  // Previous word

    if (i % N_k === 0) {
      temp = subWord(rotWord(temp))  // RotWord + SubWord
      temp = temp.map((t, index) => t ^ RCons[i / N_k][index])  // XOR with Rcon
    } else if (N_k > 6 && i % N_k === 4) {
      temp = subWord(temp)  // Apply SubWord for AES-256
    }

    // XOR with the word N_k positions earlier
    roundKeys.push(roundKeys[i - N_k].map((w, index) => w ^ temp[index]))
  }

  // Step 3: Convert flat array into 4x4 round key matrices
  const expandedKeys: number[][][] = []
  for (let i = 0;i < totalKeys / 4;i++) {
    const matrix: number[][] = [
      [roundKeys[i * 4][0], roundKeys[i * 4 + 1][0], roundKeys[i * 4 + 2][0], roundKeys[i * 4 + 3][0]],
      [roundKeys[i * 4][1], roundKeys[i * 4 + 1][1], roundKeys[i * 4 + 2][1], roundKeys[i * 4 + 3][1]],
      [roundKeys[i * 4][2], roundKeys[i * 4 + 1][2], roundKeys[i * 4 + 2][2], roundKeys[i * 4 + 3][2]],
      [roundKeys[i * 4][3], roundKeys[i * 4 + 1][3], roundKeys[i * 4 + 2][3], roundKeys[i * 4 + 3][3]],
    ]
    expandedKeys.push(matrix)
  }

  return expandedKeys
}