/*
 * GSP vGPU-IO KUnit Tests
 * In-kernel unit tests for the driver
 */

#include <kunit/test.h>
#include <kunit/mock.h>
#include <linux/slab.h>
#include <linux/string.h>
#include <linux/errno.h>
#include <linux/atomic.h>
#include <linux/list.h>
#include <linux/idr.h>
#include <linux/mutex.h>
#include <linux/kref.h>
#include <linux/completion.h>
#include <linux/wait.h>
#include <linux/spinlock.h>

#include <gsp_core.h>

/* ============================================================
 * Test: Structure Sizes and Alignments
 * ============================================================ */

static void test_structure_sizes(struct kunit *test)
{
    KUNIT_EXPECT_EQ(test, sizeof(struct gsp_rpc_header), 64);
    KUNIT_EXPECT_EQ(test, sizeof(struct gsp_rpc_message), 256);
    KUNIT_EXPECT_EQ(test, sizeof(struct gsp_rpc_response), 32);
    KUNIT_EXPECT_EQ(test, sizeof(struct gsp_subcmd_desc), 8);
    KUNIT_EXPECT_EQ(test, sizeof(struct gsp_vf_config), 64);
    KUNIT_EXPECT_EQ(test, sizeof(struct gsp_alloc_payload), 48);
    KUNIT_EXPECT_EQ(test, sizeof(struct gsp_alloc_response), 24);
    KUNIT_EXPECT_EQ(test, sizeof(struct gsp_pin_payload), 40);
    KUNIT_EXPECT_EQ(test, sizeof(struct gsp_pin_response), 8);
    KUNIT_EXPECT_EQ(test, sizeof(struct gsp_bar_trap_payload), 40);
    KUNIT_EXPECT_EQ(test, sizeof(struct gsp_bar_trap_response), 8);
    KUNIT_EXPECT_EQ(test, sizeof(struct gsp_migrate_payload), 64);
}

static void test_structure_alignments(struct kunit *test)
{
    KUNIT_EXPECT_EQ(test, __alignof__(struct gsp_rpc_header), 8);
    KUNIT_EXPECT_EQ(test, __alignof__(struct gsp_rpc_message), 8);
    KUNIT_EXPECT_EQ(test, __alignof__(struct gsp_rpc_response), 8);
    KUNIT_EXPECT_EQ(test, __alignof__(struct gsp_subcmd_desc), 4);
    KUNIT_EXPECT_EQ(test, __alignof__(struct gsp_vf_config), 8);
    KUNIT_EXPECT_EQ(test, __alignof__(struct gsp_alloc_payload), 8);
    KUNIT_EXPECT_EQ(test, __alignof__(struct gsp_alloc_response), 8);
    KUNIT_EXPECT_EQ(test, __alignof__(struct gsp_pin_payload), 8);
    KUNIT_EXPECT_EQ(test, __alignof__(struct gsp_pin_response), 8);
    KUNIT_EXPECT_EQ(test, __alignof__(struct gsp_bar_trap_payload), 8);
    KUNIT_EXPECT_EQ(test, __alignof__(struct gsp_bar_trap_response), 8);
    KUNIT_EXPECT_EQ(test, __alignof__(struct gsp_migrate_payload), 8);
}

/* ============================================================
 * Test: Payload Encoding/Decoding
 * ============================================================ */

static void test_alloc_payload_encoding(struct kunit *test)
{
    struct gsp_alloc_payload payload = {0};

    payload.size = cpu_to_le64(0x100000);
    payload.type = cpu_to_le32(GSP_MEM_VRAM);
    payload.alignment = cpu_to_le32(0x1000);
    payload.va_hint = cpu_to_le64(0x0);
    payload.flags = cpu_to_le32(0x3);
    payload.pid = cpu_to_le32(0x456);
    payload.vf_id = cpu_to_le32(0x1);

    KUNIT_EXPECT_EQ(test, le64_to_cpu(payload.size), 0x100000);
    KUNIT_EXPECT_EQ(test, le32_to_cpu(payload.type), GSP_MEM_VRAM);
    KUNIT_EXPECT_EQ(test, le32_to_cpu(payload.alignment), 0x1000);
    KUNIT_EXPECT_EQ(test, le64_to_cpu(payload.va_hint), 0x0);
    KUNIT_EXPECT_EQ(test, le32_to_cpu(payload.flags), 0x3);
    KUNIT_EXPECT_EQ(test, le32_to_cpu(payload.pid), 0x456);
    KUNIT_EXPECT_EQ(test, le32_to_cpu(payload.vf_id), 0x1);
}

static void test_pin_payload_encoding(struct kunit *test)
{
    struct gsp_pin_payload payload = {0};

    payload.handle = cpu_to_le64(0xCAFEBABE);
    payload.offset = cpu_to_le64(0x1000);
    payload.size = cpu_to_le64(0x10000);
    payload.vf_id = cpu_to_le32(0x2);
    payload.flags = cpu_to_le32(GSP_PIN_READ | GSP_PIN_WRITE);

    KUNIT_EXPECT_EQ(test, le64_to_cpu(payload.handle), 0xCAFEBABE);
    KUNIT_EXPECT_EQ(test, le64_to_cpu(payload.offset), 0x1000);
    KUNIT_EXPECT_EQ(test, le64_to_cpu(payload.size), 0x10000);
    KUNIT_EXPECT_EQ(test, le32_to_cpu(payload.vf_id), 0x2);
    KUNIT_EXPECT_EQ(test, le32_to_cpu(payload.flags), GSP_PIN_READ | GSP_PIN_WRITE);
}

static void test_bar_trap_payload_encoding(struct kunit *test)
{
    struct gsp_bar_trap_payload payload = {0};

    payload.bar_index = cpu_to_le32(1);
    payload.trap_type = cpu_to_le32(GSP_BAR_WRITE);
    payload.offset = cpu_to_le64(0x1000);
    payload.value = cpu_to_le64(0xDEADBEEF);
    payload.size = cpu_to_le32(4);
    payload.vf_id = cpu_to_le32(0x1);

    KUNIT_EXPECT_EQ(test, le32_to_cpu(payload.bar_index), 1);
    KUNIT_EXPECT_EQ(test, le32_to_cpu(payload.trap_type), GSP_BAR_WRITE);
    KUNIT_EXPECT_EQ(test, le64_to_cpu(payload.offset), 0x1000);
    KUNIT_EXPECT_EQ(test, le64_to_cpu(payload.value), 0xDEADBEEF);
    KUNIT_EXPECT_EQ(test, le32_to_cpu(payload.size), 4);
    KUNIT_EXPECT_EQ(test, le32_to_cpu(payload.vf_id), 0x1);
}

static void test_migrate_payload_encoding(struct kunit *test)
{
    struct gsp_migrate_payload payload = {0};

    payload.phase = cpu_to_le32(GSP_MIG_PREPARE);
    payload.subcmd = cpu_to_le32(3);
    payload.src_vf = cpu_to_le64(0x1);
    payload.dst_vf = cpu_to_le64(0x2);
    payload.data_ptr = cpu_to_le64(0x80000000);
    payload.data_size = cpu_to_le64(0x1000);
    payload.flags = cpu_to_le32(GSP_MIG_LIVE | GSP_MIG_VERIFY);
    payload.token = cpu_to_le64(0xABCDEF1234567890);

    KUNIT_EXPECT_EQ(test, le32_to_cpu(payload.phase), GSP_MIG_PREPARE);
    KUNIT_EXPECT_EQ(test, le32_to_cpu(payload.subcmd), 3);
    KUNIT_EXPECT_EQ(test, le64_to_cpu(payload.src_vf), 0x1);
    KUNIT_EXPECT_EQ(test, le64_to_cpu(payload.dst_vf), 0x2);
    KUNIT_EXPECT_EQ(test, le64_to_cpu(payload.data_ptr), 0x80000000);
    KUNIT_EXPECT_EQ(test, le64_to_cpu(payload.data_size), 0x1000);
    KUNIT_EXPECT_EQ(test, le32_to_cpu(payload.flags), GSP_MIG_LIVE | GSP_MIG_VERIFY);
    KUNIT_EXPECT_EQ(test, le64_to_cpu(payload.token), 0xABCDEF1234567890);
}

static void test_subcmd_desc_encoding(struct kunit *test)
{
    struct gsp_subcmd_desc desc = {0};

    desc.subcmd_id = cpu_to_le32(0x14);
    desc.echo_flags_b = 0xAA;
    desc.echo_flags_a = 0xBB;
    desc.echo_flags_c = 0xCC;

    KUNIT_EXPECT_EQ(test, le32_to_cpu(desc.subcmd_id), 0x14);
    KUNIT_EXPECT_EQ(test, desc.echo_flags_b, 0xAA);
    KUNIT_EXPECT_EQ(test, desc.echo_flags_a, 0xBB);
    KUNIT_EXPECT_EQ(test, desc.echo_flags_c, 0xCC);
}

/* ============================================================
 * Test: Firmware Address Constants
 * ============================================================ */

static void test_firmware_addresses(struct kunit *test)
{
    KUNIT_EXPECT_EQ(test, GSP_FW_BASE_ADDR, 0x04000000ULL);
    KUNIT_EXPECT_EQ(test, GSP_ENTRY_POINT, 0x04009000ULL);
    KUNIT_EXPECT_EQ(test, GSP_PRIMARY_DISPATCHER, 0x057A1AC0ULL);
    KUNIT_EXPECT_EQ(test, GSP_SECONDARY_DISPATCHER, 0x04ED39DCULL);
    KUNIT_EXPECT_EQ(test, GSP_RPC_JUMP_TABLE, 0x059BFEB0ULL);
    KUNIT_EXPECT_EQ(test, 0x04179ED8ULL, 0x04179ED8ULL);  /* Migration Table 1 */
    KUNIT_EXPECT_EQ(test, 0x04179FBCULL, 0x04179FBCULL);  /* Migration Table 2 */
    KUNIT_EXPECT_EQ(test, 0x04CB8188ULL, 0x04CB8188ULL);  /* BAR Handler */
    KUNIT_EXPECT_EQ(test, 0x04F00B0CULL, 0x04F00B0CULL);  /* WPR Handler 1 */
    KUNIT_EXPECT_EQ(test, 0x04F00C78ULL, 0x04F00C78ULL);  /* WPR Handler 2 */
    KUNIT_EXPECT_EQ(test, 0x057E94ECULL, 0x057E94ECULL);  /* ALLOC Main */
    KUNIT_EXPECT_EQ(test, 0x051EFBF4ULL, 0x051EFBF4ULL);  /* SCHEDULE Main */
}

/* ============================================================
 * Test: RPC Type Constants
 * ============================================================ */

static void test_rpc_types(struct kunit *test)
{
    KUNIT_EXPECT_EQ(test, GSP_RPC_ALLOC, 0x01);
    KUNIT_EXPECT_EQ(test, GSP_RPC_FREE, 0x02);
    KUNIT_EXPECT_EQ(test, GSP_RPC_MAP, 0x03);
    KUNIT_EXPECT_EQ(test, GSP_RPC_UNMAP, 0x04);
    KUNIT_EXPECT_EQ(test, GSP_RPC_PIN, 0x05);
    KUNIT_EXPECT_EQ(test, GSP_RPC_UNPIN, 0x06);
    KUNIT_EXPECT_EQ(test, GSP_RPC_SCHEDULE, 0x07);
    KUNIT_EXPECT_EQ(test, GSP_RPC_YIELD, 0x08);
    KUNIT_EXPECT_EQ(test, GSP_RPC_NOTIFY, 0x09);
    KUNIT_EXPECT_EQ(test, GSP_RPC_BAR_TRAP, 0x0A);
    KUNIT_EXPECT_EQ(test, GSP_RPC_MIGRATE, 0x0B);
    KUNIT_EXPECT_EQ(test, GSP_RPC_DEBUG, 0x0C);
}

static void test_migration_constants(struct kunit *test)
{
    KUNIT_EXPECT_EQ(test, GSP_MIG_QUERY, 0);
    KUNIT_EXPECT_EQ(test, GSP_MIG_PREPARE, 1);
    KUNIT_EXPECT_EQ(test, GSP_MIG_TRANSFER, 2);
    KUNIT_EXPECT_EQ(test, GSP_MIG_COMMIT, 3);
    KUNIT_EXPECT_EQ(test, GSP_MIG_CLEANUP, 4);

    KUNIT_EXPECT_EQ(test, GSP_MIG_LIVE, (1 << 0));
    KUNIT_EXPECT_EQ(test, GSP_MIG_POSTCOPY, (1 << 1));
    KUNIT_EXPECT_EQ(test, GSP_MIG_AUTO_COMMIT, (1 << 2));
    KUNIT_EXPECT_EQ(test, GSP_MIG_VERIFY, (1 << 3));
}

static void test_bar_trap_types(struct kunit *test)
{
    KUNIT_EXPECT_EQ(test, GSP_BAR_READ, 0);
    KUNIT_EXPECT_EQ(test, GSP_BAR_WRITE, 1);
    KUNIT_EXPECT_EQ(test, GSP_BAR_ATOMIC, 2);
}

static void test_pin_flags(struct kunit *test)
{
    KUNIT_EXPECT_EQ(test, GSP_PIN_READ, (1 << 0));
    KUNIT_EXPECT_EQ(test, GSP_PIN_WRITE, (1 << 1));
    KUNIT_EXPECT_EQ(test, GSP_PIN_LONGTERM, (1 << 2));
}

static void test_wpr_flags(struct kunit *test)
{
    KUNIT_EXPECT_EQ(test, GSP_WPR_WRITE, (1 << 0));
    KUNIT_EXPECT_EQ(test, GSP_WPR_READ, (1 << 1));
    KUNIT_EXPECT_EQ(test, GSP_WPR_EXECUTE, (1 << 2));
}

static void test_vf_states(struct kunit *test)
{
    KUNIT_EXPECT_EQ(test, GSP_VF_CREATING, 0);
    KUNIT_EXPECT_EQ(test, GSP_VF_ACTIVE, 1);
    KUNIT_EXPECT_EQ(test, GSP_VF_MIGRATING, 2);
    KUNIT_EXPECT_EQ(test, GSP_VF_SUSPENDED, 3);
    KUNIT_EXPECT_EQ(test, GSP_VF_DESTROYING, 4);
}

static void test_msg_flags(struct kunit *test)
{
    KUNIT_EXPECT_EQ(test, GSP_MSG_SYNC, (1 << 0));
    KUNIT_EXPECT_EQ(test, GSP_MSG_ASYNC, (1 << 1));
    KUNIT_EXPECT_EQ(test, GSP_MSG_NOREPLY, (1 << 2));
}

/* ============================================================
 * Test: Error Code Translation
 * ============================================================ */

static void test_error_codes(struct kunit *test)
{
    KUNIT_EXPECT_EQ(test, gsp_errno_to_errno(GSP_OK), 0);
    KUNIT_EXPECT_EQ(test, gsp_errno_to_errno(GSP_EPERM), -EPERM);
    KUNIT_EXPECT_EQ(test, gsp_errno_to_errno(GSP_ENOENT), -ENOENT);
    KUNIT_EXPECT_EQ(test, gsp_errno_to_errno(GSP_EINVAL), -EINVAL);
    KUNIT_EXPECT_EQ(test, gsp_errno_to_errno(GSP_ENODEV), -ENODEV);
    KUNIT_EXPECT_EQ(test, gsp_errno_to_errno(GSP_ENOMEM), -ENOMEM);
    KUNIT_EXPECT_EQ(test, gsp_errno_to_errno(GSP_EBUSY), -EBUSY);
    KUNIT_EXPECT_EQ(test, gsp_errno_to_errno(GSP_ENOSPC), -ENOSPC);
    KUNIT_EXPECT_EQ(test, gsp_errno_to_errno(GSP_ETIMEDOUT), -ETIMEDOUT);
    KUNIT_EXPECT_EQ(test, gsp_errno_to_errno(GSP_EOVERFLOW), -EOVERFLOW);
    KUNIT_EXPECT_EQ(test, gsp_errno_to_errno(GSP_EDEV), -EIO);
    KUNIT_EXPECT_EQ(test, gsp_errno_to_errno(0xDEADBEEF), -EIO);
}

/* ============================================================
 * Test: Memory Types
 * ============================================================ */

static void test_memory_types(struct kunit *test)
{
    KUNIT_EXPECT_EQ(test, GSP_MEM_VRAM, 0);
    KUNIT_EXPECT_EQ(test, GSP_MEM_SYSMEM, 1);
    KUNIT_EXPECT_EQ(test, GSP_MEM_UNCACHED, 2);
}

/* ============================================================
 * Test: Payload Capacity
 * ============================================================ */

static void test_payload_capacity(struct kunit *test)
{
    KUNIT_EXPECT_TRUE(test, sizeof(struct gsp_alloc_payload) <= 192);
    KUNIT_EXPECT_TRUE(test, sizeof(struct gsp_pin_payload) <= 192);
    KUNIT_EXPECT_TRUE(test, sizeof(struct gsp_bar_trap_payload) <= 192);
    KUNIT_EXPECT_TRUE(test, sizeof(struct gsp_migrate_payload) <= 192);
    KUNIT_EXPECT_TRUE(test, sizeof(struct gsp_subcmd_desc) <= 192);
}

/* ============================================================
 * Test Suite Registration
 * ============================================================ */

static struct kunit_case gsp_vgpu_test_cases[] = {
    KUNIT_CASE(test_structure_sizes),
    KUNIT_CASE(test_structure_alignments),
    KUNIT_CASE(test_alloc_payload_encoding),
    KUNIT_CASE(test_pin_payload_encoding),
    KUNIT_CASE(test_bar_trap_payload_encoding),
    KUNIT_CASE(test_migrate_payload_encoding),
    KUNIT_CASE(test_subcmd_desc_encoding),
    KUNIT_CASE(test_firmware_addresses),
    KUNIT_CASE(test_rpc_types),
    KUNIT_CASE(test_migration_constants),
    KUNIT_CASE(test_bar_trap_types),
    KUNIT_CASE(test_pin_flags),
    KUNIT_CASE(test_wpr_flags),
    KUNIT_CASE(test_vf_states),
    KUNIT_CASE(test_msg_flags),
    KUNIT_CASE(test_error_codes),
    KUNIT_CASE(test_memory_types),
    KUNIT_CASE(test_payload_capacity),
    {}
};

static struct kunit_suite gsp_vgpu_test_suite = {
    .name = "gsp-vgpu-io",
    .test_cases = gsp_vgpu_test_cases,
};

kunit_test_suites(&gsp_vgpu_test_suite);

MODULE_LICENSE("GPL v2");
MODULE_AUTHOR("Open-IOV Project");
MODULE_DESCRIPTION("GSP vGPU-IO KUnit Tests");