#ifndef __AES_H
#define __AES_H
#include "cstdint"
#include <string.h>

// AES-128 密钥长度 (16 bytes)
#define AES128_KEY_LEN 16
// AES 块大小 (16 bytes)
#define AES_BLOCK_SIZE 16

// AES 状态结构
typedef struct {
	u8 key[AES128_KEY_LEN];	  // 128位/16字节密钥
	u8 round_keys[11][16];	  // 10轮 + 初始轮密钥
	u8 state[AES_BLOCK_SIZE]; // 状态矩阵
} aes128_context;

// 内部函数声明（需要在头文件中声明，因为解密函数会使用）
void add_round_key(u8 *state, const u8 *round_key);
void sub_bytes(u8 *state);
void shift_rows(u8 *state);
void mix_columns(u8 *state);
void inv_sub_bytes(u8 *state);
void inv_shift_rows(u8 *state);
void inv_mix_columns(u8 *state);

// 公共函数声明
void aes128_set_key(aes128_context *ctx, const u8 *key);
void aes128_encrypt_block(aes128_context *ctx, const u8 *input, u8 *output);
void aes128_ecb_encrypt_48bytes(aes128_context *ctx, const u8 *input,
								u8 *output);
void aes128_decrypt_block(aes128_context *ctx, const u8 *input, u8 *output);
void aes128_ecb_decrypt_48bytes(aes128_context *ctx, const u8 *input,
								u8 *output);
u8 aes128_verify_encryption(aes128_context *ctx, const u8 *original_data,
							size_t length);

// 逆S盒 (用于解密)
extern const u8 inv_sbox[256];

#endif /* __AES_H */