/**
 * Copyright (c) 2026 Dominic Masters
 * 
 * This software is released under the MIT License.
 * https://opensource.org/licenses/MIT
 */

#include "dusktest.h"
#include "util/memory.h"

static void test_memoryAllocate(void **state) {
  (void)state;

  // Memory should allocate
  void *ptr = memoryAllocate(128);
  assert_non_null(ptr);

  // Allocating should be tracked
  assert_int_equal(memoryGetAllocatedCount(), 1);

  void *ptr2 = memoryAllocate(256);
  assert_non_null(ptr2);
  assert_int_equal(memoryGetAllocatedCount(), 2);

  // Free memory
  memoryFree(ptr);
  assert_int_equal(memoryGetAllocatedCount(), 1);
  memoryFree(ptr2);
  assert_int_equal(memoryGetAllocatedCount(), 0);

  // Should not be able to allocate 0 bytes
  expect_assert_failure(memoryAllocate(0));

  // Should not be able to allocate more memory than possible
  expect_assert_failure(memoryAllocate(SIZE_MAX));

  // Expect no leaks
  assert_int_equal(memoryGetAllocatedCount(), 0);
}

static void test_memoryFree(void **state) {
  (void)state;

  // Create some memory
  void *ptr = memoryAllocate(64);
  assert_non_null(ptr);
  assert_int_equal(memoryGetAllocatedCount(), 1);

  // Free
  memoryFree(ptr);

  // Allocated count should decrease
  assert_int_equal(memoryGetAllocatedCount(), 0);

  ptr = memoryAllocate(32);
  assert_non_null(ptr);
  void *ptr2 = memoryAllocate(32);
  assert_non_null(ptr2);
  assert_int_equal(memoryGetAllocatedCount(), 2);

  // Free both
  memoryFree(ptr);
  assert_int_equal(memoryGetAllocatedCount(), 1);
  memoryFree(ptr2);
  assert_int_equal(memoryGetAllocatedCount(), 0);

  // Expect unable to free NULL
  expect_assert_failure(memoryFree(NULL));

  // Expect no leaks
  assert_int_equal(memoryGetAllocatedCount(), 0);
}

static void test_memoryCopy(void **state) {
  (void)state;

  // Create some memory
  size_t size = 32;
  void *src = memoryAllocate(size);
  void *dest = memoryAllocate(size);
  assert_non_null(src);
  assert_non_null(dest);

  // Fill source with known pattern and copy
  memorySet(src, 0xAB, size);
  memoryCopy(dest, src, size);
  assert_memory_equal(src, dest, size);

  // Should not be able to copy to NULL
  expect_assert_failure(memoryCopy(NULL, src, size));

  // Should not be able to copy from NULL
  expect_assert_failure(memoryCopy(dest, NULL, size));

  // Cannot copy 0 bytes
  expect_assert_failure(memoryCopy(dest, src, 0));

  // Cannot copy to itself
  expect_assert_failure(memoryCopy(src, src, size));

  // Overlapping regions (should assert)
  void *overlap = memoryAllocate(size * 2);
  expect_assert_failure(memoryCopy((uint8_t*)overlap, (uint8_t*)overlap + 8, size));
  expect_assert_failure(memoryCopy((uint8_t*)overlap + 8, (uint8_t*)overlap, size));
  memoryFree(overlap);

  memoryFree(src);
  memoryFree(dest);

  // Expect no leaks
  assert_int_equal(memoryGetAllocatedCount(), 0);
}

static void test_memorySet(void **state) {
  (void)state;

  // Allocate memory
  size_t size = 16;
  void *ptr = memoryAllocate(size);
  assert_non_null(ptr);

  // Fill with pattern
  memorySet(ptr, 0xCD, size);
  uint8_t expected[16];
  for(size_t i = 0; i < size; i++) {
    expected[i] = 0xCD;
  }
  assert_memory_equal(ptr, expected, size);

  // Cannot set to NULL
  expect_assert_failure(memorySet(NULL, 0x00, size));

  // Canot set 0 bytes
  expect_assert_failure(memorySet(ptr, 0x00, 0));

  // Set with value outside 0-255 (should be cast to uint8_t)
  memorySet(ptr, 0x1FF, size);
  for(size_t i = 0; i < size; i++) {
    assert_int_equal(((uint8_t*)ptr)[i], 0xFF);
  }

  memoryFree(ptr);

  // Expect no leaks
  assert_int_equal(memoryGetAllocatedCount(), 0);
}

static void test_memoryZero(void **state) {
  (void)state;

  // Create some memory
  size_t size = 20;
  void *ptr = memoryAllocate(size);
  assert_non_null(ptr);

  // Fill and then zero
  memorySet(ptr, 0xEF, size);
  memoryZero(ptr, size);

  // All memory should be zeroed
  uint8_t expected[20] = {0};
  assert_memory_equal(ptr, expected, size);

  // Cannot zero to NULL pointer
  expect_assert_failure(memoryZero(NULL, size));

  // Cannot zero 0 bytes
  expect_assert_failure(memoryZero(ptr, 0));

  memoryFree(ptr);

  // Expect no leaks
  assert_int_equal(memoryGetAllocatedCount(), 0);
}

static void test_memoryCopyRangeSafe(void **state) {
  (void)state;

  // Create some memory
  size_t size = 10;
  void *src = memoryAllocate(size);
  void *dest = memoryAllocate(size);
  assert_non_null(src);
  assert_non_null(dest);

  // Initialize source with known pattern
  for(size_t i = 0; i < size; i++) {
    ((uint8_t*)src)[i] = (uint8_t)(i + 1); // 1, 2, ..., 10
  }

  // Copy a range from src[2] to src[7] (5 bytes) into dest
  memoryCopyRangeSafe(dest, (uint8_t*)src + 2, (uint8_t*)src + 7, size);
  uint8_t expected[10] = {3, 4, 5, 6, 7}; // Expected data in dest
  assert_memory_equal(dest, expected, 5);

  // Copy entire buffer
  memoryCopyRangeSafe(dest, src, (uint8_t*)src + size, size);
  assert_memory_equal(dest, src, size);

  // Cannot copy to NULL destination
  expect_assert_failure(memoryCopyRangeSafe(NULL, src, (uint8_t*)src + 5, size));
  
  // Cannot copy from NULL start
  expect_assert_failure(memoryCopyRangeSafe(dest, NULL, (uint8_t*)src + 5, size));
  
  // Cannot copy from NULL end
  expect_assert_failure(memoryCopyRangeSafe(dest, (uint8_t*)src, NULL, size));
  
  // Start and end are the same
  expect_assert_failure(memoryCopyRangeSafe(dest, (uint8_t*)src, (uint8_t*)src, size));
  
  // End is before start
  expect_assert_failure(memoryCopyRangeSafe(dest, (uint8_t*)src + 5, (uint8_t*)src + 2, size));
  
  // Size to copy exceeds maximum
  expect_assert_failure(memoryCopyRangeSafe(dest, (uint8_t*)src, (uint8_t*)src + 10, 5));

  memoryFree(src);
  memoryFree(dest);

  // Expect no leaks
  assert_int_equal(memoryGetAllocatedCount(), 0);
}

static void test_memoryMove(void **state) {
  (void)state;

  // Create some memory
  size_t size = 15;
  void *ptr = memoryAllocate(size + 5); // Extra space for overlap
  assert_non_null(ptr);

  // Initialize memory with known pattern
  for(size_t i = 0; i < size + 5; i++) {
    ((uint8_t*)ptr)[i] = (uint8_t)(i + 1); // 1, 2, ..., size+5
  }

  // Move overlapping region: from ptr[0..14] to ptr[5..19]
  memoryMove((uint8_t*)ptr + 5, ptr, size);

  // Expected pattern after move
  uint8_t expected[20];
  for(size_t i = 0; i < 5; i++) {
    expected[i] = (uint8_t)(i + 1); // Original first 5 bytes
  }
  for(size_t i = 5; i < 20; i++) {
    expected[i] = (uint8_t)(i - 4); // Moved bytes
  }
  assert_memory_equal(ptr, expected, size + 5);

  // Overlap backward: move ptr[5..19] to ptr[0..14]
  for(size_t i = 0; i < size + 5; i++) {
    ((uint8_t*)ptr)[i] = (uint8_t)(i + 1);
  }
  memoryMove(ptr, (uint8_t*)ptr + 5, size);
  for(size_t i = 0; i < size; i++) {
    assert_int_equal(((uint8_t*)ptr)[i], (uint8_t)(i + 6));
  }

  // Cannot move to NULL
  expect_assert_failure(memoryMove(NULL, ptr, size));

  // Cannot move from NULL
  expect_assert_failure(memoryMove(ptr, NULL, size));

  // Cannot move 0 bytes
  expect_assert_failure(memoryMove(ptr, ptr, 0));

  // Cannot move to itself
  expect_assert_failure(memoryMove(ptr, ptr, size));

  memoryFree(ptr);

  // Expect no leaks
  assert_int_equal(memoryGetAllocatedCount(), 0);
}

static void test_memoryCompare(void **state) {
  (void)state;

  // Create two memory blocks
  size_t size = 8;
  void *a = memoryAllocate(size);
  void *b = memoryAllocate(size);
  assert_non_null(a);
  assert_non_null(b);

  // Initialize both with same data
  for(size_t i = 0; i < size; i++) {
    ((uint8_t*)a)[i] = (uint8_t)(i + 1);
    ((uint8_t*)b)[i] = (uint8_t)(i + 1);
  }

  // They should be equal
  assert_int_equal(memoryCompare(a, b, size), 0);

  // Change b and compare again
  ((uint8_t*)b)[size - 1] = 0xFF;
  assert_true(memoryCompare(a, b, size) < 0);

  // Change a and compare again
  ((uint8_t*)a)[size - 1] = 0xFF;
  ((uint8_t*)a)[size - 2] = 0xFE;
  assert_true(memoryCompare(a, b, size) > 0);

  // CAnnot compare with NULL a
  expect_assert_failure(memoryCompare(NULL, b, size));

  // Cannot compare with NULL b
  expect_assert_failure(memoryCompare(a, NULL, size));

  // Comparing with self always equal
  assert_int_equal(memoryCompare(a, a, size), 0);
  assert_int_equal(memoryCompare(b, b, size), 0);

  // Cannot compare 0 bytes
  expect_assert_failure(memoryCompare(a, b, 0));

  // Compare different sizes (should assert, so we only test the API contract)
  // Not possible with current API, as size is explicit, but document intent.

  memoryFree(a);
  memoryFree(b);

  // Expect no leaks
  assert_int_equal(memoryGetAllocatedCount(), 0);
}

static void test_memoryReallocate(void **state) {
  (void)state;

  size_t initialSize = 16;
  void *ptr = memoryAllocate(initialSize);
  assert_non_null(ptr);

  // Reallocate to a larger size
  size_t newSize = 32;
  memoryReallocate(&ptr, newSize);
  assert_non_null(ptr);

  // Reallocate to a smaller size
  size_t smallerSize = 8;
  memoryReallocate(&ptr, smallerSize);
  assert_non_null(ptr);

  // Cannot realloc to size 0
  expect_assert_failure(memoryReallocate(&ptr, 0));

  // Cannot realloc NULL pointer
  expect_assert_failure(memoryReallocate(NULL, 16));

  // Cannot reallocate more memory than possible
  expect_assert_failure(memoryReallocate(&ptr, SIZE_MAX));

  // All we really care about is that the pointer is valid after reallocations
  memoryFree(ptr);

  // Expect no leaks
  assert_int_equal(memoryGetAllocatedCount(), 0);
}

static void test_memoryResize(void **state) {
  (void)state;

  size_t initialSize = 32;
  void *ptr = memoryAllocate(initialSize);
  assert_non_null(ptr);

  // Initialize memory
  for(size_t i = 0; i < initialSize; i++) {
    ((uint8_t*)ptr)[i] = (uint8_t)(i + 1);
  }

  // Equal size shouldn't touch the pointer
  void *oldPtr = ptr;
  memoryResize(&ptr, initialSize, initialSize);
  assert_non_null(ptr);
  assert_ptr_equal(ptr, oldPtr);

  // Reallocate to a larger size
  size_t newSize = 64;
  memoryResize(&ptr, initialSize, newSize);
  assert_non_null(ptr);

  // Check that old data is preserved
  for(size_t i = 0; i < initialSize; i++) {
    assert_int_equal(((uint8_t*)ptr)[i], (uint8_t)(i + 1));
  }

  // Reallocate to a smaller size
  size_t smallerSize = 16;
  memoryResize(&ptr, newSize, smallerSize);
  assert_non_null(ptr);

  // Check that data is still correct up to the smaller size
  for(size_t i = 0; i < smallerSize; i++) {
    assert_int_equal(((uint8_t*)ptr)[i], (uint8_t)(i + 1));
  }

  // Cannot realloc to size 0
  expect_assert_failure(memoryResize(&ptr, smallerSize, 0));

  // Cannot take NULL
  expect_assert_failure(memoryResize(NULL, smallerSize, 16));

  // memoryResize with oldsize of 0 not possible
  expect_assert_failure(memoryResize(&ptr, 0, 16));
  
  // Cannot resize more memory than possible
  expect_assert_failure(memoryResize(&ptr, smallerSize, SIZE_MAX));

  memoryFree(ptr);
  
  // Expect no leaks
  assert_int_equal(memoryGetAllocatedCount(), 0);
}

int main(int argc, char **argv) {
  const struct CMUnitTest tests[] = {
    cmocka_unit_test(test_memoryAllocate),
    cmocka_unit_test(test_memoryFree),
    cmocka_unit_test(test_memoryCopy),
    cmocka_unit_test(test_memorySet),
    cmocka_unit_test(test_memoryZero),
    cmocka_unit_test(test_memoryCopyRangeSafe),
    cmocka_unit_test(test_memoryMove),
    cmocka_unit_test(test_memoryCompare),
    cmocka_unit_test(test_memoryReallocate),
    cmocka_unit_test(test_memoryResize),
  };
  return cmocka_run_group_tests(tests, NULL, NULL);
}