// SPDX-License-Identifier: Apache-2.0
// Copyright 2025 AES Encryption Project
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#ifndef AES_ENCRYPTION_H
#define AES_ENCRYPTION_H

#include <vector>
#include <string>
#include <cstdint>

/**
 * @brief AES (Advanced Encryption Standard) implementation
 * 
 * This class provides AES-128, AES-192, and AES-256 encryption and decryption
 * capabilities using Electronic Codebook (ECB) and Cipher Block Chaining (CBC) modes.
 */
class AesEncryption {
public:
    // Lifecycle
    AesEncryption();
    ~AesEncryption();
    
    // Key management
    bool SetKey(const std::vector<uint8_t>& key);
    bool SetKey(const std::string& hexKey);
    void ClearKey();
    
    // Encryption/Decryption operations
    std::vector<uint8_t> EncryptEcb(const std::vector<uint8_t>& plaintext);
    std::vector<uint8_t> DecryptEcb(const std::vector<uint8_t>& ciphertext);
    
    std::vector<uint8_t> EncryptCbc(const std::vector<uint8_t>& plaintext, 
                                   const std::vector<uint8_t>& iv);
    std::vector<uint8_t> DecryptCbc(const std::vector<uint8_t>& ciphertext, 
                                   const std::vector<uint8_t>& iv);
    
    // Utility methods
    std::string ToHexString(const std::vector<uint8_t>& data);
    std::vector<uint8_t> FromHexString(const std::string& hexStr);
    std::vector<uint8_t> GenerateRandomIv();
    
    // Validation
    bool IsKeySet() const;
    int GetKeySize() const;
    
private:
    // AES constants
    static const int BLOCK_SIZE = 16;
    static const int MAX_ROUNDS = 14;
    static const int MAX_KEY_SIZE = 32;
    
    // Internal state
    std::vector<uint8_t> mKey;
    std::vector<std::vector<uint8_t>> mRoundKeys;
    int mKeySize;
    int mRounds;
    bool mKeyIsSet;
    
    // Core AES operations
    void KeyExpansion();
    void EncryptBlock(const uint8_t* input, uint8_t* output);
    void DecryptBlock(const uint8_t* input, uint8_t* output);
    
    // AES transformation functions
    void SubBytes(uint8_t state[4][4]);
    void InvSubBytes(uint8_t state[4][4]);
    void ShiftRows(uint8_t state[4][4]);
    void InvShiftRows(uint8_t state[4][4]);
    void MixColumns(uint8_t state[4][4]);
    void InvMixColumns(uint8_t state[4][4]);
    void AddRoundKey(uint8_t state[4][4], int round);
    
    // Helper functions
    uint8_t GfMultiply(uint8_t a, uint8_t b);
    uint8_t SubByte(uint8_t byte);
    uint8_t InvSubByte(uint8_t byte);
    void RotWord(uint8_t* word);
    void SubWord(uint8_t* word);
    uint8_t Rcon(int i);
    
    // Padding operations
    std::vector<uint8_t> AddPkcs7Padding(const std::vector<uint8_t>& data);
    std::vector<uint8_t> RemovePkcs7Padding(const std::vector<uint8_t>& data);
    
    // State conversion utilities
    void BytesToState(const uint8_t* bytes, uint8_t state[4][4]);
    void StateToBytes(const uint8_t state[4][4], uint8_t* bytes);
    
    // S-box and inverse S-box tables
    static const uint8_t S_BOX[256];
    static const uint8_t INV_S_BOX[256];
};

#endif // AES_ENCRYPTION_H