#pragma once
#include <vector>
#include <array>
#include <cstdint>
#include <cstring>
#include <stdexcept>
#include <algorithm>

namespace netplus {

using block128 = std::vector<uint8_t>;
using key128   = std::vector<uint8_t>;

class aes {
protected:
    // GCM operations - implemented by subclasses
    virtual void gcm_mul(std::array<uint8_t,16>& X) const = 0;
    virtual void ghash(std::array<uint8_t,16>& out,
                       const uint8_t* aad, size_t aad_len,
                       const uint8_t* ct,  size_t ct_len) const = 0;

    // Allow aes_gcm to access protected methods
    friend class aes_gcm;

public:
    virtual ~aes() = default;
    virtual std::vector<uint8_t> encrypt(const std::vector<uint8_t>& plaintext) = 0;
    virtual std::vector<uint8_t> decrypt(const std::vector<uint8_t>& ciphertext) = 0;

    virtual std::vector<uint8_t> encryptCBC(const std::vector<uint8_t>& plaintext,
                                            const std::vector<uint8_t>& iv) = 0;
    virtual std::vector<uint8_t> decryptCBC(const std::vector<uint8_t>& ciphertext,
                                            const std::vector<uint8_t>& iv) = 0;

    // GCM encryption/decryption
    virtual bool aes_gcm_encrypt(
        const uint8_t iv12[12],
        const uint8_t* aad, size_t aad_len,
        const uint8_t* pt,  size_t pt_len,
        uint8_t* ct,
        uint8_t tag16[16]
    ) const = 0;

    virtual bool aes_gcm_decrypt(
        const uint8_t iv12[12],
        const uint8_t* aad, size_t aad_len,
        const uint8_t* ct,  size_t ct_len,
        const uint8_t tag16[16],
        uint8_t* pt
    ) const = 0;
};

class aes128 : public aes {
private:
    std::vector<uint8_t> m_roundKeyBytes; // 176 bytes for AES-128
    uint8_t m_H[16];                      // GHASH subkey

    inline const uint8_t* rk(int round) const {
        return &m_roundKeyBytes[size_t(round) * 16];
    }

    // ---- Constant-time helper ops ----
    static inline uint8_t ct_eq_u8(uint8_t a, uint8_t b) {
        uint8_t x = static_cast<uint8_t>(a ^ b);
        x |= x >> 4; x |= x >> 2; x |= x >> 1;
        return static_cast<uint8_t>((x ^ 1) & 1);
    }

    static inline uint8_t ct_select_u8(uint8_t mask, uint8_t a, uint8_t b) {
        // mask: 0x00 or 0xFF
        return static_cast<uint8_t>((a & mask) | (b & ~mask));
    }

    static inline uint8_t rotl8(uint8_t x, unsigned r) {
        return static_cast<uint8_t>((x << r) | (x >> (8 - r)));
    }

    // ---- GF(2^8) multiply for AES S-box ----
    static uint8_t gf_mul(uint8_t a, uint8_t b);
    static uint8_t gf_inv(uint8_t a);

    // ---- constant-time AES S-box + inv S-box ----
    static uint8_t ct_sbox(uint8_t x);
    static uint8_t ct_inv_sbox(uint8_t x);

    // ---- key schedule ----
    static uint32_t rotWord(uint32_t w);
    static uint32_t subWord_ct(uint32_t w);
    void keyExpansion(const key128& key);

    // ---- AES core block ops ----
    void encrypt_block(const uint8_t in[16], uint8_t out[16]) const;
    void decrypt_block(const uint8_t in[16], uint8_t out[16]) const;

    static void addRoundKey(uint8_t* state, const uint8_t* roundKey);
    static void subBytes(uint8_t* state);
    static void invSubBytes(uint8_t* state);
    static void shiftRows(uint8_t* state);
    static void invShiftRows(uint8_t* state);
    static void mixColumns(uint8_t* state);
    static void invMixColumns(uint8_t* state);

    static inline uint8_t xtime(uint8_t x) {
        return static_cast<uint8_t>((x << 1) ^ ((x & 0x80) ? 0x1B : 0x00));
    }

    static uint8_t gmul(uint8_t a, uint8_t b);

    // ---- GCM helpers ----
    static inline void inc32(uint8_t ctr[16]) {
        uint32_t x = (uint32_t(ctr[12]) << 24) |
                    (uint32_t(ctr[13]) << 16) |
                    (uint32_t(ctr[14]) << 8)  |
                    (uint32_t(ctr[15]));
        x++;
        ctr[12] = uint8_t(x >> 24);
        ctr[13] = uint8_t(x >> 16);
        ctr[14] = uint8_t(x >> 8);
        ctr[15] = uint8_t(x);
    }

    static inline uint64_t load_be64(const uint8_t* p) {
        return (uint64_t(p[0]) << 56) |
               (uint64_t(p[1]) << 48) |
               (uint64_t(p[2]) << 40) |
               (uint64_t(p[3]) << 32) |
               (uint64_t(p[4]) << 24) |
               (uint64_t(p[5]) << 16) |
               (uint64_t(p[6]) << 8)  |
               (uint64_t(p[7]));
    }

    static inline void store_be64(uint8_t* p, uint64_t x) {
        p[0] = uint8_t(x >> 56);
        p[1] = uint8_t(x >> 48);
        p[2] = uint8_t(x >> 40);
        p[3] = uint8_t(x >> 32);
        p[4] = uint8_t(x >> 24);
        p[5] = uint8_t(x >> 16);
        p[6] = uint8_t(x >> 8);
        p[7] = uint8_t(x);
    }

    static inline bool ct_memeq(const uint8_t* a, const uint8_t* b, size_t n) {
        uint8_t diff = 0;
        for (size_t i = 0; i < n; i++)
            diff |= static_cast<uint8_t>(a[i] ^ b[i]);
        return diff == 0;
    }

protected:
    void gcm_mul(std::array<uint8_t,16>& X) const override;
    void ghash(std::array<uint8_t,16>& out,
               const uint8_t* aad, size_t aad_len,
               const uint8_t* ct,  size_t ct_len) const override;

public:
    explicit aes128(const key128& key);
    ~aes128() = default;

    std::vector<uint8_t> encrypt(const std::vector<uint8_t>& plaintext) override;
    std::vector<uint8_t> decrypt(const std::vector<uint8_t>& ciphertext) override;

    std::vector<uint8_t> encryptCBC(const std::vector<uint8_t>& plaintext,
                                    const std::vector<uint8_t>& iv) override;

    std::vector<uint8_t> decryptCBC(const std::vector<uint8_t>& ciphertext,
                                    const std::vector<uint8_t>& iv) override;

    // CTS mode (RFC 3962 - Kerberos AES encryption)
    // Input may be any length >= 16 bytes (one AES block).
    // No padding is needed; output size equals input size.
    std::vector<uint8_t> encryptCTS(const std::vector<uint8_t>& plaintext,
                                    const std::vector<uint8_t>& iv);

    std::vector<uint8_t> decryptCTS(const std::vector<uint8_t>& ciphertext,
                                    const std::vector<uint8_t>& iv);

    bool aes_gcm_encrypt(
        const uint8_t iv12[12],
        const uint8_t* aad, size_t aad_len,
        const uint8_t* pt,  size_t pt_len,
        uint8_t* ct,
        uint8_t tag16[16]
    ) const  override;

    bool aes_gcm_decrypt(
        const uint8_t iv12[12],
        const uint8_t* aad, size_t aad_len,
        const uint8_t* ct,  size_t ct_len,
        const uint8_t tag16[16],
        uint8_t* pt
    ) const override;
};

class aes256 : public aes {
private:
    std::vector<uint8_t> m_roundKeyBytes; // 240 bytes for AES-256
    uint8_t m_H[16];                      // GHASH subkey

    inline const uint8_t* rk(int round) const {
        return &m_roundKeyBytes[size_t(round) * 16];
    }

    static uint8_t gf_mul(uint8_t a, uint8_t b);
    static uint8_t gf_inv(uint8_t a);
    static uint8_t ct_sbox(uint8_t x);
    static uint8_t ct_inv_sbox(uint8_t x);

    static uint32_t rotWord(uint32_t w);
    static uint32_t subWord_ct(uint32_t w);
    void keyExpansion(const std::vector<uint8_t>& key);

    void encrypt_block(const uint8_t in[16], uint8_t out[16]) const;
    void decrypt_block(const uint8_t in[16], uint8_t out[16]) const;

    static void addRoundKey(uint8_t* state, const uint8_t* roundKey);
    static void subBytes(uint8_t* state);
    static void invSubBytes(uint8_t* state);
    static void shiftRows(uint8_t* state);
    static void invShiftRows(uint8_t* state);
    static void mixColumns(uint8_t* state);
    static void invMixColumns(uint8_t* state);

    static inline uint8_t xtime(uint8_t x) {
        return static_cast<uint8_t>((x << 1) ^ ((x & 0x80) ? 0x1B : 0x00));
    }

    static uint8_t gmul(uint8_t a, uint8_t b);

    static inline void inc32(uint8_t ctr[16]) {
        uint32_t x = (uint32_t(ctr[12]) << 24) |
                    (uint32_t(ctr[13]) << 16) |
                    (uint32_t(ctr[14]) << 8)  |
                    (uint32_t(ctr[15]));
        x++;
        ctr[12] = uint8_t(x >> 24);
        ctr[13] = uint8_t(x >> 16);
        ctr[14] = uint8_t(x >> 8);
        ctr[15] = uint8_t(x);
    }

    static inline uint64_t load_be64(const uint8_t* p) {
        return (uint64_t(p[0]) << 56) |
               (uint64_t(p[1]) << 48) |
               (uint64_t(p[2]) << 40) |
               (uint64_t(p[3]) << 32) |
               (uint64_t(p[4]) << 24) |
               (uint64_t(p[5]) << 16) |
               (uint64_t(p[6]) << 8)  |
               (uint64_t(p[7]));
    }

    static inline void store_be64(uint8_t* p, uint64_t x) {
        p[0] = uint8_t(x >> 56);
        p[1] = uint8_t(x >> 48);
        p[2] = uint8_t(x >> 40);
        p[3] = uint8_t(x >> 32);
        p[4] = uint8_t(x >> 24);
        p[5] = uint8_t(x >> 16);
        p[6] = uint8_t(x >> 8);
        p[7] = uint8_t(x);
    }

    static inline bool ct_memeq(const uint8_t* a, const uint8_t* b, size_t n) {
        uint8_t diff = 0;
        for (size_t i = 0; i < n; i++)
            diff |= static_cast<uint8_t>(a[i] ^ b[i]);
        return diff == 0;
    }

protected:
    void gcm_mul(std::array<uint8_t,16>& X) const override;
    void ghash(std::array<uint8_t,16>& out,
               const uint8_t* aad, size_t aad_len,
               const uint8_t* ct,  size_t ct_len) const override;

public:
    explicit aes256(const std::vector<uint8_t>& key);
    ~aes256() = default;

    std::vector<uint8_t> encrypt(const std::vector<uint8_t>& plaintext) override;
    std::vector<uint8_t> decrypt(const std::vector<uint8_t>& ciphertext) override;

    std::vector<uint8_t> encryptCBC(const std::vector<uint8_t>& plaintext,
                                    const std::vector<uint8_t>& iv) override;

    std::vector<uint8_t> decryptCBC(const std::vector<uint8_t>& ciphertext,
                                    const std::vector<uint8_t>& iv) override;

    // CTS mode (RFC 3962 - Kerberos AES encryption)
    std::vector<uint8_t> encryptCTS(const std::vector<uint8_t>& plaintext,
                                    const std::vector<uint8_t>& iv);

    std::vector<uint8_t> decryptCTS(const std::vector<uint8_t>& ciphertext,
                                    const std::vector<uint8_t>& iv);

    bool aes_gcm_encrypt(
        const uint8_t iv12[12],
        const uint8_t* aad, size_t aad_len,
        const uint8_t* pt,  size_t pt_len,
        uint8_t* ct,
        uint8_t tag16[16]
    ) const override;

    bool aes_gcm_decrypt(
        const uint8_t iv12[12],
        const uint8_t* aad, size_t aad_len,
        const uint8_t* ct,  size_t ct_len,
        const uint8_t tag16[16],
        uint8_t* pt
    ) const override;
};

class aes_gcm {
public:
    static bool encrypt(
        const aes& cipher,
        const uint8_t iv12[12],
        const uint8_t* aad, size_t aad_len,
        const uint8_t* pt,  size_t pt_len,
        uint8_t* ct,
        uint8_t tag16[16]
    );

    static bool decrypt(
        const aes& cipher,
        const uint8_t iv12[12],
        const uint8_t* aad, size_t aad_len,
        const uint8_t* ct,  size_t ct_len,
        const uint8_t tag16[16],
        uint8_t* pt
    );
};


} // namespace netplus

