#include "../src/common/memory.h"
#include "munit.h"
#include <stdint.h>
#include <string.h>

MunitResult test_alloc_success(const MunitParameter params[], void *data) {
    (void)params;
    (void)data;

    void *ptr = nullptr;
    error_t *err;

    // Test various allocation sizes
    err = mem_alloc(&ptr, 1, 1);
    munit_assert_null(err);
    munit_assert_not_null(ptr);
    mem_free(&ptr);
    munit_assert_null(ptr);

    err = mem_alloc(&ptr, 10, sizeof(int));
    munit_assert_null(err);
    munit_assert_not_null(ptr);
    mem_free(&ptr);

    err = mem_alloc(&ptr, 1024, 1);
    munit_assert_null(err);
    munit_assert_not_null(ptr);
    mem_free(&ptr);

    err = mem_alloc(&ptr, 100, 100);
    munit_assert_null(err);
    munit_assert_not_null(ptr);
    mem_free(&ptr);

    return MUNIT_OK;
}

MunitResult test_alloc_zeroed(const MunitParameter params[], void *data) {
    (void)params;
    (void)data;

    void *ptr = nullptr;
    error_t *err = mem_alloc(&ptr, 256, sizeof(uint8_t));
    munit_assert_null(err);
    munit_assert_not_null(ptr);

    // Verify all bytes are zero-initialized
    uint8_t *bytes = (uint8_t *)ptr;
    for (size_t i = 0; i < 256; i++) {
        munit_assert_uint8(bytes[i], ==, 0);
    }

    mem_free(&ptr);
    return MUNIT_OK;
}

MunitResult test_alloc_invalid_null_ptr(const MunitParameter params[], void *data) {
    (void)params;
    (void)data;

    error_t *err = mem_alloc(nullptr, 10, sizeof(int));
    munit_assert_ptr_equal(err, err_invalid_parameters);

    return MUNIT_OK;
}

MunitResult test_alloc_invalid_zero_n(const MunitParameter params[], void *data) {
    (void)params;
    (void)data;

    void *ptr = nullptr;
    error_t *err = mem_alloc(&ptr, 0, sizeof(int));
    munit_assert_ptr_equal(err, err_invalid_parameters);
    munit_assert_null(ptr);

    return MUNIT_OK;
}

MunitResult test_alloc_invalid_zero_size(const MunitParameter params[], void *data) {
    (void)params;
    (void)data;

    void *ptr = nullptr;
    error_t *err = mem_alloc(&ptr, 10, 0);
    munit_assert_ptr_equal(err, err_invalid_parameters);
    munit_assert_null(ptr);

    return MUNIT_OK;
}

MunitResult test_alloc_invalid_zero_both(const MunitParameter params[], void *data) {
    (void)params;
    (void)data;

    void *ptr = nullptr;
    error_t *err = mem_alloc(&ptr, 0, 0);
    munit_assert_ptr_equal(err, err_invalid_parameters);
    munit_assert_null(ptr);

    return MUNIT_OK;
}

MunitResult test_alloc_integer_overflow(const MunitParameter params[], void *data) {
    (void)params;
    (void)data;

    void *ptr = nullptr;
    error_t *err = mem_alloc(&ptr, SIZE_MAX, SIZE_MAX);
    munit_assert_ptr_equal(err, err_integer_overflow);
    munit_assert_null(ptr);

    // Test case where multiplication would overflow
    err = mem_alloc(&ptr, SIZE_MAX / 2 + 1, 2);
    munit_assert_ptr_equal(err, err_integer_overflow);
    munit_assert_null(ptr);

    return MUNIT_OK;
}

MunitResult test_resize_grow(const MunitParameter params[], void *data) {
    (void)params;
    (void)data;

    void *ptr = nullptr;
    error_t *err;

    // Allocate initial memory
    err = mem_alloc(&ptr, 10, sizeof(uint32_t));
    munit_assert_null(err);
    munit_assert_not_null(ptr);

    // Write a pattern to the initial allocation
    uint32_t *ints = (uint32_t *)ptr;
    for (size_t i = 0; i < 10; i++) {
        ints[i] = 0xDEADBEEF + i;
    }

    // Resize to larger size
    err = mem_resize(&ptr, 20, sizeof(uint32_t));
    munit_assert_null(err);
    munit_assert_not_null(ptr);

    // Verify original data is preserved
    ints = (uint32_t *)ptr;
    for (size_t i = 0; i < 10; i++) {
        munit_assert_uint32(ints[i], ==, 0xDEADBEEF + i);
    }

    // Verify new bytes are zero-initialized
    uint8_t *bytes = (uint8_t *)ptr;
    for (size_t i = 10 * sizeof(uint32_t); i < 20 * sizeof(uint32_t); i++) {
        munit_assert_uint8(bytes[i], ==, 0);
    }

    mem_free(&ptr);
    return MUNIT_OK;
}

MunitResult test_resize_shrink(const MunitParameter params[], void *data) {
    (void)params;
    (void)data;

    void *ptr = nullptr;
    error_t *err;

    // Allocate initial memory
    err = mem_alloc(&ptr, 20, sizeof(uint32_t));
    munit_assert_null(err);
    munit_assert_not_null(ptr);

    // Write a pattern to the initial allocation
    uint32_t *ints = (uint32_t *)ptr;
    for (size_t i = 0; i < 20; i++) {
        ints[i] = 0xCAFEBABE + i;
    }

    // Resize to smaller size
    err = mem_resize(&ptr, 10, sizeof(uint32_t));
    munit_assert_null(err);
    munit_assert_not_null(ptr);

    // Verify preserved data is intact
    ints = (uint32_t *)ptr;
    for (size_t i = 0; i < 10; i++) {
        munit_assert_uint32(ints[i], ==, 0xCAFEBABE + i);
    }

    mem_free(&ptr);
    return MUNIT_OK;
}

MunitResult test_resize_same_size(const MunitParameter params[], void *data) {
    (void)params;
    (void)data;

    void *ptr = nullptr;
    error_t *err;

    // Allocate memory
    err = mem_alloc(&ptr, 15, sizeof(uint32_t));
    munit_assert_null(err);
    munit_assert_not_null(ptr);

    // Write a pattern
    uint32_t *ints = (uint32_t *)ptr;
    for (size_t i = 0; i < 15; i++) {
        ints[i] = 0x12345678 + i;
    }

    // Resize to the same size
    err = mem_resize(&ptr, 15, sizeof(uint32_t));
    munit_assert_null(err);
    munit_assert_not_null(ptr);

    // Verify data is unchanged
    ints = (uint32_t *)ptr;
    for (size_t i = 0; i < 15; i++) {
        munit_assert_uint32(ints[i], ==, 0x12345678 + i);
    }

    mem_free(&ptr);
    return MUNIT_OK;
}

MunitResult test_resize_invalid_params(const MunitParameter params[], void *data) {
    (void)params;
    (void)data;

    void *ptr = nullptr;
    error_t *err;

    // Test null pointer parameter
    err = mem_resize(nullptr, 10, sizeof(int));
    munit_assert_ptr_equal(err, err_invalid_parameters);

    // Test valid ptr but *ptr is nullptr
    err = mem_resize(&ptr, 10, sizeof(int));
    munit_assert_ptr_equal(err, err_invalid_parameters);

    // Allocate memory for other tests
    err = mem_alloc(&ptr, 10, sizeof(int));
    munit_assert_null(err);

    // Test zero n
    err = mem_resize(&ptr, 0, sizeof(int));
    munit_assert_ptr_equal(err, err_invalid_parameters);

    // Test zero size
    err = mem_resize(&ptr, 10, 0);
    munit_assert_ptr_equal(err, err_invalid_parameters);

    // Test integer overflow
    err = mem_resize(&ptr, SIZE_MAX, SIZE_MAX);
    munit_assert_ptr_equal(err, err_integer_overflow);

    // Test case where multiplication would overflow
    err = mem_resize(&ptr, SIZE_MAX / 2 + 1, 2);
    munit_assert_ptr_equal(err, err_integer_overflow);

    mem_free(&ptr);
    return MUNIT_OK;
}

MunitResult test_resize_failure_preserves_ptr(const MunitParameter params[], void *data) {
    (void)params;
    (void)data;

    void *ptr = nullptr;
    void *original_ptr;
    error_t *err;

    // Allocate initial memory
    err = mem_alloc(&ptr, 10, sizeof(uint32_t));
    munit_assert_null(err);
    munit_assert_not_null(ptr);

    // Write a pattern to verify pointer validity
    uint32_t *ints = (uint32_t *)ptr;
    for (size_t i = 0; i < 10; i++) {
        ints[i] = 0xBEEFCAFE + i;
    }

    // Save original pointer
    original_ptr = ptr;

    // Attempt resize with overflow - should fail
    err = mem_resize(&ptr, SIZE_MAX, SIZE_MAX);
    munit_assert_ptr_equal(err, err_integer_overflow);

    // Verify ptr is unchanged
    munit_assert_ptr_equal(ptr, original_ptr);

    // Verify data is still intact
    ints = (uint32_t *)ptr;
    for (size_t i = 0; i < 10; i++) {
        munit_assert_uint32(ints[i], ==, 0xBEEFCAFE + i);
    }

    // Verify we can still operate on the pointer (resize successfully)
    err = mem_resize(&ptr, 20, sizeof(uint32_t));
    munit_assert_null(err);
    munit_assert_not_null(ptr);

    // Verify original data preserved after successful resize
    ints = (uint32_t *)ptr;
    for (size_t i = 0; i < 10; i++) {
        munit_assert_uint32(ints[i], ==, 0xBEEFCAFE + i);
    }

    // Verify we can free the pointer
    mem_free(&ptr);
    munit_assert_null(ptr);

    return MUNIT_OK;
}

MunitResult test_resize_minimal_to_minimal(const MunitParameter params[], void *data) {
    (void)params;
    (void)data;

    void *ptr = nullptr;
    error_t *err;

    // Allocate minimal size
    err = mem_alloc(&ptr, 1, 1);
    munit_assert_null(err);
    munit_assert_not_null(ptr);

    // Write a byte
    uint8_t *byte = (uint8_t *)ptr;
    *byte = 0x42;

    // Resize to same minimal size
    err = mem_resize(&ptr, 1, 1);
    munit_assert_null(err);
    munit_assert_not_null(ptr);

    // Verify data is preserved
    byte = (uint8_t *)ptr;
    munit_assert_uint8(*byte, ==, 0x42);

    mem_free(&ptr);
    return MUNIT_OK;
}

MunitResult test_resize_shrink_then_grow_zeroinit(const MunitParameter params[], void *data) {
    (void)params;
    (void)data;

    void *ptr = nullptr;
    error_t *err;

    // Allocate 100 bytes
    err = mem_alloc(&ptr, 100, 1);
    munit_assert_null(err);
    munit_assert_not_null(ptr);

    // Write pattern throughout
    uint8_t *bytes = (uint8_t *)ptr;
    for (size_t i = 0; i < 100; i++) {
        bytes[i] = (uint8_t)(0xAA + (i % 16));
    }

    // Shrink to 50 bytes
    err = mem_resize(&ptr, 50, 1);
    munit_assert_null(err);
    munit_assert_not_null(ptr);

    // Verify first 50 bytes preserved
    bytes = (uint8_t *)ptr;
    for (size_t i = 0; i < 50; i++) {
        munit_assert_uint8(bytes[i], ==, (uint8_t)(0xAA + (i % 16)));
    }

    // Grow to 150 bytes
    err = mem_resize(&ptr, 150, 1);
    munit_assert_null(err);
    munit_assert_not_null(ptr);

    // Verify first 50 bytes still preserved
    bytes = (uint8_t *)ptr;
    for (size_t i = 0; i < 50; i++) {
        munit_assert_uint8(bytes[i], ==, (uint8_t)(0xAA + (i % 16)));
    }

    // Verify bytes 50-149 are zero-initialized
    for (size_t i = 50; i < 150; i++) {
        munit_assert_uint8(bytes[i], ==, 0);
    }

    mem_free(&ptr);
    return MUNIT_OK;
}

MunitResult test_resize_odd_boundaries_zeroinit(const MunitParameter params[], void *data) {
    (void)params;
    (void)data;

    void *ptr = nullptr;
    error_t *err;

    // Allocate odd-sized 13 bytes
    err = mem_alloc(&ptr, 13, 1);
    munit_assert_null(err);
    munit_assert_not_null(ptr);

    // Write pattern to all 13 bytes
    uint8_t *bytes = (uint8_t *)ptr;
    for (size_t i = 0; i < 13; i++) {
        bytes[i] = (uint8_t)(0xCC + i);
    }

    // Resize to odd-sized 27 bytes
    err = mem_resize(&ptr, 27, 1);
    munit_assert_null(err);
    munit_assert_not_null(ptr);

    // Verify original 13 bytes preserved
    bytes = (uint8_t *)ptr;
    for (size_t i = 0; i < 13; i++) {
        munit_assert_uint8(bytes[i], ==, (uint8_t)(0xCC + i));
    }

    // Verify new 14 bytes (13-26) are zero-initialized
    for (size_t i = 13; i < 27; i++) {
        munit_assert_uint8(bytes[i], ==, 0);
    }

    mem_free(&ptr);
    return MUNIT_OK;
}

MunitResult test_resize_multiple_times(const MunitParameter params[], void *data) {
    (void)params;
    (void)data;

    void *ptr = nullptr;
    error_t *err;

    // Start with small allocation
    err = mem_alloc(&ptr, 5, sizeof(uint32_t));
    munit_assert_null(err);
    munit_assert_not_null(ptr);

    // Write initial pattern
    uint32_t *ints = (uint32_t *)ptr;
    for (size_t i = 0; i < 5; i++) {
        ints[i] = 0xAAAA0000 + i;
    }

    // Resize 1: grow to 10
    err = mem_resize(&ptr, 10, sizeof(uint32_t));
    munit_assert_null(err);
    ints = (uint32_t *)ptr;
    for (size_t i = 0; i < 5; i++) {
        munit_assert_uint32(ints[i], ==, 0xAAAA0000 + i);
    }

    // Resize 2: grow to 20
    err = mem_resize(&ptr, 20, sizeof(uint32_t));
    munit_assert_null(err);
    ints = (uint32_t *)ptr;
    for (size_t i = 0; i < 5; i++) {
        munit_assert_uint32(ints[i], ==, 0xAAAA0000 + i);
    }

    // Resize 3: shrink to 15
    err = mem_resize(&ptr, 15, sizeof(uint32_t));
    munit_assert_null(err);
    ints = (uint32_t *)ptr;
    for (size_t i = 0; i < 5; i++) {
        munit_assert_uint32(ints[i], ==, 0xAAAA0000 + i);
    }

    // Resize 4: shrink to 3
    err = mem_resize(&ptr, 3, sizeof(uint32_t));
    munit_assert_null(err);
    ints = (uint32_t *)ptr;
    for (size_t i = 0; i < 3; i++) {
        munit_assert_uint32(ints[i], ==, 0xAAAA0000 + i);
    }

    // Resize 5: grow back to 25
    err = mem_resize(&ptr, 25, sizeof(uint32_t));
    munit_assert_null(err);
    ints = (uint32_t *)ptr;
    for (size_t i = 0; i < 3; i++) {
        munit_assert_uint32(ints[i], ==, 0xAAAA0000 + i);
    }

    mem_free(&ptr);
    return MUNIT_OK;
}

MunitResult test_free_normal(const MunitParameter params[], void *data) {
    (void)params;
    (void)data;

    void *ptr = nullptr;
    error_t *err;

    // Allocate memory
    err = mem_alloc(&ptr, 50, sizeof(uint32_t));
    munit_assert_null(err);
    munit_assert_not_null(ptr);

    // Free it
    mem_free(&ptr);
    munit_assert_null(ptr);

    return MUNIT_OK;
}

MunitResult test_free_null_ptr_param(const MunitParameter params[], void *data) {
    (void)params;
    (void)data;

    // Freeing with nullptr parameter should be a no-op (not crash)
    mem_free(nullptr);

    return MUNIT_OK;
}

MunitResult test_free_null_deref_ptr(const MunitParameter params[], void *data) {
    (void)params;
    (void)data;

    void *ptr = nullptr;

    // Freeing a pointer that points to nullptr should be a no-op
    mem_free(&ptr);
    munit_assert_null(ptr);

    return MUNIT_OK;
}

MunitResult test_free_double_free(const MunitParameter params[], void *data) {
    (void)params;
    (void)data;

    void *ptr = nullptr;
    error_t *err;

    // Allocate memory
    err = mem_alloc(&ptr, 10, sizeof(int));
    munit_assert_null(err);
    munit_assert_not_null(ptr);

    // Free once
    mem_free(&ptr);
    munit_assert_null(ptr);

    // Free again - should be a no-op since ptr is now nullptr
    mem_free(&ptr);
    munit_assert_null(ptr);

    return MUNIT_OK;
}

MunitResult test_free_after_resize(const MunitParameter params[], void *data) {
    (void)params;
    (void)data;

    void *ptr = nullptr;
    error_t *err;

    // Allocate memory
    err = mem_alloc(&ptr, 5, sizeof(uint64_t));
    munit_assert_null(err);
    munit_assert_not_null(ptr);

    // Resize it
    err = mem_resize(&ptr, 20, sizeof(uint64_t));
    munit_assert_null(err);
    munit_assert_not_null(ptr);

    // Free the resized memory
    mem_free(&ptr);
    munit_assert_null(ptr);

    return MUNIT_OK;
}

MunitResult test_free_multiple_allocations(const MunitParameter params[], void *data) {
    (void)params;
    (void)data;

    void *ptr1 = nullptr, *ptr2 = nullptr, *ptr3 = nullptr;
    error_t *err;

    // Allocate multiple blocks
    err = mem_alloc(&ptr1, 10, sizeof(int));
    munit_assert_null(err);
    munit_assert_not_null(ptr1);

    err = mem_alloc(&ptr2, 20, sizeof(char));
    munit_assert_null(err);
    munit_assert_not_null(ptr2);

    err = mem_alloc(&ptr3, 30, sizeof(double));
    munit_assert_null(err);
    munit_assert_not_null(ptr3);

    // Free them in different order than allocated
    mem_free(&ptr2);
    munit_assert_null(ptr2);

    mem_free(&ptr1);
    munit_assert_null(ptr1);

    mem_free(&ptr3);
    munit_assert_null(ptr3);

    return MUNIT_OK;
}

MunitTest memory_tests[] = {
    {.name = "/alloc/success", .test = test_alloc_success},
    {.name = "/alloc/zeroed", .test = test_alloc_zeroed},
    {.name = "/alloc/invalid_null_ptr", .test = test_alloc_invalid_null_ptr},
    {.name = "/alloc/invalid_zero_n", .test = test_alloc_invalid_zero_n},
    {.name = "/alloc/invalid_zero_size", .test = test_alloc_invalid_zero_size},
    {.name = "/alloc/invalid_zero_both", .test = test_alloc_invalid_zero_both},
    {.name = "/alloc/integer_overflow", .test = test_alloc_integer_overflow},
    {.name = "/resize/grow", .test = test_resize_grow},
    {.name = "/resize/shrink", .test = test_resize_shrink},
    {.name = "/resize/same_size", .test = test_resize_same_size},
    {.name = "/resize/invalid_params", .test = test_resize_invalid_params},
    {.name = "/resize/failure_preserves_ptr", .test = test_resize_failure_preserves_ptr},
    {.name = "/resize/minimal_to_minimal", .test = test_resize_minimal_to_minimal},
    {.name = "/resize/shrink_then_grow_zeroinit", .test = test_resize_shrink_then_grow_zeroinit},
    {.name = "/resize/odd_boundaries_zeroinit", .test = test_resize_odd_boundaries_zeroinit},
    {.name = "/resize/multiple_times", .test = test_resize_multiple_times},
    {.name = "/free/normal", .test = test_free_normal},
    {.name = "/free/null_ptr_param", .test = test_free_null_ptr_param},
    {.name = "/free/null_deref_ptr", .test = test_free_null_deref_ptr},
    {.name = "/free/double_free", .test = test_free_double_free},
    {.name = "/free/after_resize", .test = test_free_after_resize},
    {.name = "/free/multiple_allocations", .test = test_free_multiple_allocations},
    {}
};