/**
 * @file test_scp03_key_rotation.c
 * @brief Platform SCP03 Key Rotation Test Suite
 * License: MIT (Clean-room implementation)
 */

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdint.h>
#include <fcntl.h>
#include <unistd.h>
#include <sys/ioctl.h>
#include <linux/i2c.h>
#include <linux/i2c-dev.h>

#include "se050_wireguard.h"
#include "se050_crypto_utils.h"
#include "se050_scp03_keys.h"

#ifndef SE050_CHIP
#define SE050_CHIP 1
#endif

#define CHIP_SE050C0    0
#define CHIP_SE050C1    1
#define CHIP_SE050C2    2
#define CHIP_SE050E2    3

#if SE050_CHIP == CHIP_SE050C1
    #define CHIP_NAME "SE050C1"
    #define SE050_DEFAULT_I2C_ADDR 0x48
    #define DEFAULT_ENC_KEY SE050C1_ENC_KEY
    #define DEFAULT_MAC_KEY SE050C1_MAC_KEY
    #define DEFAULT_DEK_KEY SE050C1_DEK_KEY
#elif SE050_CHIP == CHIP_SE050C2
    #define CHIP_NAME "SE050C2"
    #define SE050_DEFAULT_I2C_ADDR 0x48
    #define DEFAULT_ENC_KEY SE050C2_ENC_KEY
    #define DEFAULT_MAC_KEY SE050C2_MAC_KEY
    #define DEFAULT_DEK_KEY SE050C2_DEK_KEY
#elif SE050_CHIP == CHIP_SE050E2
    #define CHIP_NAME "SE050E2"
    #define SE050_DEFAULT_I2C_ADDR 0x48
    #define DEFAULT_ENC_KEY SE050E2_ENC_KEY
    #define DEFAULT_MAC_KEY SE050E2_MAC_KEY
    #define DEFAULT_DEK_KEY SE050E2_DEK_KEY
#else
    #error "Invalid SE050_CHIP. Use SE050C1, SE050C2, or SE050E2"
#endif

static const uint8_t TEST_ENC_KEY[16] = {
    0xAA, 0xBB, 0xCC, 0xDD, 0xEE, 0xFF, 0x00, 0x11,
    0x22, 0x33, 0x44, 0x55, 0x66, 0x77, 0x88, 0x99
};

static const uint8_t TEST_MAC_KEY[16] = {
    0x11, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77, 0x88,
    0x99, 0xAA, 0xBB, 0xCC, 0xDD, 0xEE, 0xFF, 0x00
};

static const uint8_t TEST_DEK_KEY[16] = {
    0x0F, 0x1E, 0x2D, 0x3C, 0x4B, 0x5A, 0x69, 0x78,
    0x87, 0x96, 0xA5, 0xB4, 0xC3, 0xD2, 0xE1, 0xF0
};

static int test_passed = 0;
static int test_failed = 0;

#define TEST_ASSERT_EQ(a, b, msg) do { if ((a) == (b)) { printf("[PASS] %s\n", msg); test_passed++; } else { printf("[FAIL] %s\n", msg); test_failed++; } } while(0)

typedef struct { int fd; uint8_t slave_addr; const char *dev_path; } i2c_ctx_t;

static int i2c_init(i2c_ctx_t *ctx, const char *dev_path, uint8_t addr) {
    int fd = open(dev_path, O_RDWR);
    if (fd < 0) return -1;
    if (ioctl(fd, I2C_SLAVE, addr) < 0) { close(fd); return -1; }
    ctx->fd = fd; ctx->slave_addr = addr; ctx->dev_path = dev_path;
    return 0;
}

static void i2c_close(i2c_ctx_t *ctx) { if (ctx->fd >= 0) { close(ctx->fd); ctx->fd = -1; } }

static int step1_open_default_keys(const char *i2c_bus, se050_session_ctx_t **session) {
    printf("\n=== Step 1: Session Open with Default Keys ===\n");
    printf("Chip: %s\n", CHIP_NAME);
    i2c_ctx_t i2c; se050_i2c_hal_t hal;
    if (i2c_init(&i2c, i2c_bus, SE050_DEFAULT_I2C_ADDR) != 0) { printf("[FAIL] I2C connection failed\n"); return -1; }
    memset(&hal, 0, sizeof(hal)); hal.handle = &i2c; hal.slave_addr = SE050_DEFAULT_I2C_ADDR; hal.dev_path = i2c_bus;
    se050_status_t status = se050_session_create(session, &hal);
    if (status != SE050_OK) { printf("[FAIL] Session creation failed (0x%04x)\n", status); i2c_close(&i2c); return -1; }
    status = se050_session_scp03_init(*session);
    if (status != SE050_OK) { se050_session_delete(*session); i2c_close(&i2c); return -1; }
    status = se050_session_scp03_set_keys(*session, DEFAULT_ENC_KEY, DEFAULT_MAC_KEY, DEFAULT_DEK_KEY);
    if (status != SE050_OK) { se050_session_delete(*session); i2c_close(&i2c); return -1; }
    printf("[PASS] Session opened with default keys\n");
    return 0;
}

static void step2_basic_tests(se050_session_ctx_t *session) {
    printf("\n=== Step 2: Basic Tests with Default Keys ===\n");
    if (!session) { printf("[FAIL] No session\n"); return; }
    uint8_t cmd[64]; size_t cmd_len = 6;
    cmd[0] = 0x80; cmd[1] = 0x6F; cmd[2] = 0x00; cmd[3] = 0x00; cmd[4] = 0x00; cmd[5] = 0x00;
    se050_status_t status = se050_session_scp03_encrypt(session, cmd, &cmd_len);
    TEST_ASSERT_EQ(status, SE050_OK, "SCP03 encryption with default keys");
}

static int step3_rotate_to_test_keys(se050_session_ctx_t *session) {
    printf("\n=== Step 3: Rotate to Test Keys ===\n");
    if (!session) { printf("[FAIL] No session\n"); return -1; }
    se050_status_t status = se050_session_scp03_set_keys(session, TEST_ENC_KEY, TEST_MAC_KEY, TEST_DEK_KEY);
    if (status != SE050_OK) { printf("[FAIL] Set test keys failed (0x%04x)\n", status); return -1; }
    printf("[PASS] Rotated to test keys\n");
    return 0;
}

static int step4_reopen_test_keys(const char *i2c_bus, se050_session_ctx_t **session) {
    printf("\n=== Step 4: Close and Reopen with Test Keys ===\n");
    if (*session) { se050_session_delete(*session); printf("[INFO] Session closed\n"); }
    i2c_ctx_t i2c; se050_i2c_hal_t hal;
    if (i2c_init(&i2c, i2c_bus, SE050_DEFAULT_I2C_ADDR) != 0) { printf("[FAIL] I2C connection failed\n"); return -1; }
    memset(&hal, 0, sizeof(hal)); hal.handle = &i2c; hal.slave_addr = SE050_DEFAULT_I2C_ADDR; hal.dev_path = i2c_bus;
    se050_status_t status = se050_session_create(session, &hal);
    if (status != SE050_OK) { printf("[FAIL] Session creation failed (0x%04x)\n", status); i2c_close(&i2c); return -1; }
    status = se050_session_scp03_init(*session);
    if (status != SE050_OK) { se050_session_delete(*session); i2c_close(&i2c); return -1; }
    status = se050_session_scp03_set_keys(*session, TEST_ENC_KEY, TEST_MAC_KEY, TEST_DEK_KEY);
    if (status != SE050_OK) { se050_session_delete(*session); i2c_close(&i2c); return -1; }
    printf("[PASS] Session reopened with test keys\n");
    return 0;
}

static int step5_rotate_back_default(se050_session_ctx_t *session) {
    printf("\n=== Step 5: Rotate Back to Default Keys ===\n");
    if (!session) { printf("[FAIL] No session\n"); return -1; }
    se050_status_t status = se050_session_scp03_set_keys(session, DEFAULT_ENC_KEY, DEFAULT_MAC_KEY, DEFAULT_DEK_KEY);
    if (status != SE050_OK) { printf("[FAIL] Set default keys failed (0x%04x)\n", status); return -1; }
    printf("[PASS] Rotated back to default keys\n");
    return 0;
}

static int step6_reopen_default_keys(const char *i2c_bus, se050_session_ctx_t **session) {
    printf("\n=== Step 6: Close and Reopen with Default Keys ===\n");
    if (*session) { se050_session_delete(*session); printf("[INFO] Session closed\n"); }
    i2c_ctx_t i2c; se050_i2c_hal_t hal;
    if (i2c_init(&i2c, i2c_bus, SE050_DEFAULT_I2C_ADDR) != 0) { printf("[FAIL] I2C connection failed\n"); return -1; }
    memset(&hal, 0, sizeof(hal)); hal.handle = &i2c; hal.slave_addr = SE050_DEFAULT_I2C_ADDR; hal.dev_path = i2c_bus;
    se050_status_t status = se050_session_create(session, &hal);
    if (status != SE050_OK) { printf("[FAIL] Session creation failed (0x%04x)\n", status); i2c_close(&i2c); return -1; }
    status = se050_session_scp03_init(*session);
    if (status != SE050_OK) { se050_session_delete(*session); i2c_close(&i2c); return -1; }
    status = se050_session_scp03_set_keys(*session, DEFAULT_ENC_KEY, DEFAULT_MAC_KEY, DEFAULT_DEK_KEY);
    if (status != SE050_OK) { se050_session_delete(*session); i2c_close(&i2c); return -1; }
    printf("[PASS] Session reopened with default keys\n");
    return 0;
}

static void step7_final_close(se050_session_ctx_t **session) {
    printf("\n=== Step 7: Final Close ===\n");
    if (*session) { se050_session_delete(*session); printf("[PASS] Session closed with default keys\n"); *session = NULL; }
    else { printf("[FAIL] No session to close\n"); }
}

int main(int argc, char *argv[]) {
    const char *i2c_bus = "/dev/i2c-1";
    for (int i = 1; i < argc; i++) { if (strcmp(argv[i], "-b") == 0 && i + 1 < argc) { i2c_bus = argv[++i]; } }
    printf("========================================\n");
    printf("Platform SCP03 Key Rotation Test\n");
    printf("========================================\n");
    printf("Chip Type: %s\n", CHIP_NAME);
    printf("I2C Bus: %s\n", i2c_bus);
    printf("========================================\n");
    se050_session_ctx_t *session = NULL; int result = 0;
    if (step1_open_default_keys(i2c_bus, &session) != 0) { printf("\n[ERROR] Step 1 failed - requires SE050 hardware\n"); return 1; }
    step2_basic_tests(session);
    if (step3_rotate_to_test_keys(session) != 0) { result = 1; goto cleanup; }
    if (step4_reopen_test_keys(i2c_bus, &session) != 0) { result = 1; goto cleanup; }
    if (step5_rotate_back_default(session) != 0) { result = 1; goto cleanup; }
    if (step6_reopen_default_keys(i2c_bus, &session) != 0) { result = 1; goto cleanup; }
    step7_final_close(&session);
cleanup:
    if (session) se050_session_delete(session);
    printf("\n========================================\n");
    printf("Test Summary\n");
    printf("Passed: %d\n", test_passed);
    printf("Failed: %d\n", test_failed);
    printf("========================================\n");
    return (test_failed > 0 || result != 0) ? 1 : 0;
}