"""Pure-Python AES block encryption.

Only the encryption direction is implemented, as that is all FAST requires.
"""

from __future__ import annotations

SBOX = bytes(
    [
        0x63,
        0x7C,
        0x77,
        0x7B,
        0xF2,
        0x6B,
        0x6F,
        0xC5,
        0x30,
        0x01,
        0x67,
        0x2B,
        0xFE,
        0xD7,
        0xAB,
        0x76,
        0xCA,
        0x82,
        0xC9,
        0x7D,
        0xFA,
        0x59,
        0x47,
        0xF0,
        0xAD,
        0xD4,
        0xA2,
        0xAF,
        0x9C,
        0xA4,
        0x72,
        0xC0,
        0xB7,
        0xFD,
        0x93,
        0x26,
        0x36,
        0x3F,
        0xF7,
        0xCC,
        0x34,
        0xA5,
        0xE5,
        0xF1,
        0x71,
        0xD8,
        0x31,
        0x15,
        0x04,
        0xC7,
        0x23,
        0xC3,
        0x18,
        0x96,
        0x05,
        0x9A,
        0x07,
        0x12,
        0x80,
        0xE2,
        0xEB,
        0x27,
        0xB2,
        0x75,
        0x09,
        0x83,
        0x2C,
        0x1A,
        0x1B,
        0x6E,
        0x5A,
        0xA0,
        0x52,
        0x3B,
        0xD6,
        0xB3,
        0x29,
        0xE3,
        0x2F,
        0x84,
        0x53,
        0xD1,
        0x00,
        0xED,
        0x20,
        0xFC,
        0xB1,
        0x5B,
        0x6A,
        0xCB,
        0xBE,
        0x39,
        0x4A,
        0x4C,
        0x58,
        0xCF,
        0xD0,
        0xEF,
        0xAA,
        0xFB,
        0x43,
        0x4D,
        0x33,
        0x85,
        0x45,
        0xF9,
        0x02,
        0x7F,
        0x50,
        0x3C,
        0x9F,
        0xA8,
        0x51,
        0xA3,
        0x40,
        0x8F,
        0x92,
        0x9D,
        0x38,
        0xF5,
        0xBC,
        0xB6,
        0xDA,
        0x21,
        0x10,
        0xFF,
        0xF3,
        0xD2,
        0xCD,
        0x0C,
        0x13,
        0xEC,
        0x5F,
        0x97,
        0x44,
        0x17,
        0xC4,
        0xA7,
        0x7E,
        0x3D,
        0x64,
        0x5D,
        0x19,
        0x73,
        0x60,
        0x81,
        0x4F,
        0xDC,
        0x22,
        0x2A,
        0x90,
        0x88,
        0x46,
        0xEE,
        0xB8,
        0x14,
        0xDE,
        0x5E,
        0x0B,
        0xDB,
        0xE0,
        0x32,
        0x3A,
        0x0A,
        0x49,
        0x06,
        0x24,
        0x5C,
        0xC2,
        0xD3,
        0xAC,
        0x62,
        0x91,
        0x95,
        0xE4,
        0x79,
        0xE7,
        0xC8,
        0x37,
        0x6D,
        0x8D,
        0xD5,
        0x4E,
        0xA9,
        0x6C,
        0x56,
        0xF4,
        0xEA,
        0x65,
        0x7A,
        0xAE,
        0x08,
        0xBA,
        0x78,
        0x25,
        0x2E,
        0x1C,
        0xA6,
        0xB4,
        0xC6,
        0xE8,
        0xDD,
        0x74,
        0x1F,
        0x4B,
        0xBD,
        0x8B,
        0x8A,
        0x70,
        0x3E,
        0xB5,
        0x66,
        0x48,
        0x03,
        0xF6,
        0x0E,
        0x61,
        0x35,
        0x57,
        0xB9,
        0x86,
        0xC1,
        0x1D,
        0x9E,
        0xE1,
        0xF8,
        0x98,
        0x11,
        0x69,
        0xD9,
        0x8E,
        0x94,
        0x9B,
        0x1E,
        0x87,
        0xE9,
        0xCE,
        0x55,
        0x28,
        0xDF,
        0x8C,
        0xA1,
        0x89,
        0x0D,
        0xBF,
        0xE6,
        0x42,
        0x68,
        0x41,
        0x99,
        0x2D,
        0x0F,
        0xB0,
        0x54,
        0xBB,
        0x16,
    ]
)

RCON = bytes([0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80, 0x1B, 0x36])

MUL2 = bytes([(x << 1) & 0xFF ^ (0x1B if x & 0x80 else 0) & 0xFF for x in range(256)])
MUL3 = bytes(
    [((x << 1) & 0xFF ^ (0x1B if x & 0x80 else 0) & 0xFF) ^ x for x in range(256)]
)


def _sub_bytes(state: bytes) -> bytes:
    return bytes(SBOX[b] for b in state)


def _xor_bytes(a: bytes | bytearray, b: bytes | bytearray) -> bytes:
    return bytes(x ^ y for x, y in zip(a, b))


def _shift_rows(state: bytes) -> bytes:
    return bytes(
        [
            state[0],
            state[5],
            state[10],
            state[15],
            state[4],
            state[9],
            state[14],
            state[3],
            state[8],
            state[13],
            state[2],
            state[7],
            state[12],
            state[1],
            state[6],
            state[11],
        ]
    )


def _mix_columns(state: bytes) -> bytes:
    new_state = bytearray(16)
    for i in range(4):
        s0, s1, s2, s3 = state[i * 4 : i * 4 + 4]
        new_state[i * 4] = MUL2[s0] ^ MUL3[s1] ^ s2 ^ s3
        new_state[i * 4 + 1] = s0 ^ MUL2[s1] ^ MUL3[s2] ^ s3
        new_state[i * 4 + 2] = s0 ^ s1 ^ MUL2[s2] ^ MUL3[s3]
        new_state[i * 4 + 3] = MUL3[s0] ^ s1 ^ s2 ^ MUL2[s3]
    return bytes(new_state)


def _expand_key(key: bytes) -> list[bytes]:
    key_len = len(key)
    if key_len not in (16, 24, 32):
        raise ValueError("Key must be 16, 24, or 32 bytes")

    nk = key_len // 4  # number of 32-bit words in key
    nr = nk + 6  # number of rounds
    total_words = 4 * (nr + 1)

    # Split key into 4-byte words
    w: list[bytes] = [key[4 * i : 4 * i + 4] for i in range(nk)]

    for i in range(nk, total_words):
        temp = w[i - 1]
        if i % nk == 0:
            temp = temp[1:] + temp[:1]
            temp = _sub_bytes(temp)
            temp = bytes([temp[0] ^ RCON[i // nk - 1]]) + temp[1:]
        elif nk > 6 and i % nk == 4:
            temp = _sub_bytes(temp)
        w.append(_xor_bytes(w[i - nk], temp))

    return [b"".join(w[4 * i : 4 * i + 4]) for i in range(nr + 1)]


def _encrypt_block(round_keys: list[bytes], block: bytes) -> bytes:
    nr = len(round_keys) - 1
    state = _xor_bytes(block, round_keys[0])
    for i in range(nr - 1):
        state = _sub_bytes(state)
        state = _shift_rows(state)
        state = _mix_columns(state)
        state = _xor_bytes(state, round_keys[i + 1])
    state = _sub_bytes(state)
    state = _shift_rows(state)
    state = _xor_bytes(state, round_keys[nr])
    return state


class AesEncryptor:
    """AES block encryptor. Supports 128/192/256-bit keys."""

    def __init__(self, key: bytes) -> None:
        if len(key) not in (16, 24, 32):
            raise ValueError("Key must be 16, 24, or 32 bytes")
        self._round_keys = _expand_key(key)

    def encrypt_block(self, block: bytes) -> bytes:
        return _encrypt_block(self._round_keys, block)
