#define DOCTEST_CONFIG_IMPLEMENT_WITH_MAIN
#include "arena_allocator.hpp"
#include <cstring>
#include <doctest/doctest.h>
#include <vector>

TEST_CASE("ArenaAllocator basic construction") {
  ArenaAllocator arena;
  CHECK(arena.num_blocks() == 0);
  CHECK(arena.used_bytes() == 0);
  CHECK(arena.total_allocated() == 0);
  CHECK(arena.available_in_current_block() == 0);
}

TEST_CASE("ArenaAllocator custom initial size") {
  ArenaAllocator arena(2048);
  CHECK(arena.num_blocks() == 0);
  CHECK(arena.total_allocated() == 0);
  CHECK(arena.available_in_current_block() == 0);
}

TEST_CASE("ArenaAllocator basic allocation") {
  ArenaAllocator arena;

  SUBCASE("allocate zero bytes returns nullptr") {
    void *ptr = arena.allocate_raw(0);
    CHECK(ptr == nullptr);
    CHECK(arena.used_bytes() == 0);
  }

  SUBCASE("allocate single byte") {
    void *ptr = arena.allocate_raw(1);
    CHECK(ptr != nullptr);
    CHECK(arena.used_bytes() >= 1);
  }

  SUBCASE("allocate multiple bytes") {
    void *ptr1 = arena.allocate_raw(100);
    void *ptr2 = arena.allocate_raw(200);

    CHECK(ptr1 != nullptr);
    CHECK(ptr2 != nullptr);
    CHECK(ptr1 != ptr2);
    CHECK(arena.used_bytes() >= 300);
  }
}

TEST_CASE("ArenaAllocator alignment") {
  ArenaAllocator arena;

  SUBCASE("default alignment") {
    void *ptr = arena.allocate_raw(1);
    CHECK(reinterpret_cast<uintptr_t>(ptr) % alignof(std::max_align_t) == 0);
  }

  SUBCASE("custom alignment") {
    void *ptr8 = arena.allocate_raw(1, 8);
    CHECK(reinterpret_cast<uintptr_t>(ptr8) % 8 == 0);

    void *ptr16 = arena.allocate_raw(1, 16);
    CHECK(reinterpret_cast<uintptr_t>(ptr16) % 16 == 0);

    void *ptr32 = arena.allocate_raw(1, 32);
    CHECK(reinterpret_cast<uintptr_t>(ptr32) % 32 == 0);
  }

  SUBCASE("alignment with larger allocations") {
    ArenaAllocator fresh_arena;
    void *ptr = fresh_arena.allocate_raw(100, 64);
    CHECK(reinterpret_cast<uintptr_t>(ptr) % 64 == 0);
  }
}

TEST_CASE("ArenaAllocator block management") {
  ArenaAllocator arena(128);

  SUBCASE("single block allocation") {
    void *ptr = arena.allocate_raw(64);
    CHECK(ptr != nullptr);
    CHECK(arena.num_blocks() == 1);
    CHECK(arena.used_bytes() == 64);
  }

  SUBCASE("multiple blocks when size exceeded") {
    void *ptr1 = arena.allocate_raw(100);
    CHECK(arena.num_blocks() == 1);

    void *ptr2 = arena.allocate_raw(50);
    CHECK(arena.num_blocks() == 2);
    CHECK(ptr1 != ptr2);
  }

  SUBCASE("allocation larger than block size grows arena") {
    void *ptr = arena.allocate_raw(200);
    CHECK(ptr != nullptr);
    CHECK(arena.num_blocks() == 1);
  }
}

TEST_CASE("ArenaAllocator construct template") {
  ArenaAllocator arena;

  SUBCASE("construct int") {
    int *ptr = arena.construct<int>(42);
    CHECK(ptr != nullptr);
    CHECK(*ptr == 42);
  }

  SUBCASE("construct array") {
    struct FixedString {
      char data[16];
      FixedString(const char *str) {
        std::strncpy(data, str, sizeof(data) - 1);
        data[sizeof(data) - 1] = '\0';
      }
    };
    FixedString *ptr = arena.construct<FixedString>("hello world");
    CHECK(ptr != nullptr);
    CHECK(std::strcmp(ptr->data, "hello world") == 0);
  }

  SUBCASE("construct multiple objects") {
    int *ptr1 = arena.construct<int>(10);
    int *ptr2 = arena.construct<int>(20);

    CHECK(ptr1 != ptr2);
    CHECK(*ptr1 == 10);
    CHECK(*ptr2 == 20);
  }

  SUBCASE("construct with multiple arguments") {
    struct IntPair {
      int first;
      int second;
      IntPair(int a, int b) : first(a), second(b) {}
    };
    auto *ptr = arena.construct<IntPair>(42, 24);
    CHECK(ptr != nullptr);
    CHECK(ptr->first == 42);
    CHECK(ptr->second == 24);
  }
}

TEST_CASE("ArenaAllocator reset functionality") {
  ArenaAllocator arena;

  arena.allocate_raw(100);
  arena.allocate_raw(200);
  size_t used_before = arena.used_bytes();
  CHECK(used_before > 0);

  arena.reset();
  CHECK(arena.used_bytes() == 0);
  CHECK(arena.num_blocks() >= 1);

  void *ptr = arena.allocate_raw(50);
  CHECK(ptr != nullptr);
  CHECK(arena.used_bytes() == 50);
}

TEST_CASE("ArenaAllocator reset memory leak test") {
  ArenaAllocator arena(32); // Smaller initial size

  // Force multiple blocks
  arena.allocate_raw(30); // First block (32 bytes)
  CHECK(arena.num_blocks() == 1);

  arena.allocate_raw(
      30); // Should create second block (64 bytes due to doubling)
  CHECK(arena.num_blocks() == 2);

  arena.allocate_raw(100); // Should create third block (128 bytes due to
                           // doubling, or larger for 100)
  CHECK(arena.num_blocks() == 3);

  size_t total_before = arena.total_allocated();
  CHECK(total_before > 32);

  arena.reset();

  // After reset, only first block should remain (others freed to prevent memory
  // leak)
  CHECK(arena.num_blocks() == 1);
  CHECK(arena.total_allocated() == 32); // Only first block size
  CHECK(arena.used_bytes() == 0);

  // Should be able to use the first block again
  void *ptr = arena.allocate_raw(20);
  CHECK(ptr != nullptr);
  CHECK(arena.used_bytes() == 20);
}

TEST_CASE("ArenaAllocator memory tracking") {
  ArenaAllocator arena(512);

  CHECK(arena.total_allocated() == 0);
  CHECK(arena.used_bytes() == 0);
  CHECK(arena.available_in_current_block() == 0);

  arena.allocate_raw(100);
  CHECK(arena.used_bytes() >= 100);
  CHECK(arena.available_in_current_block() <= 412);

  arena.allocate_raw(400);
  CHECK(arena.num_blocks() == 1);

  arena.allocate_raw(50);
  CHECK(arena.num_blocks() == 2);
  CHECK(arena.total_allocated() >= 1024);
}

TEST_CASE("ArenaAllocator stress test") {
  ArenaAllocator arena(1024);

  SUBCASE("many small allocations") {
    std::vector<void *> ptrs;
    for (int i = 0; i < 1000; ++i) {
      void *ptr = arena.allocate_raw(8);
      CHECK(ptr != nullptr);
      ptrs.push_back(ptr);
    }

    for (size_t i = 1; i < ptrs.size(); ++i) {
      CHECK(ptrs[i] != ptrs[i - 1]);
    }
  }

  SUBCASE("alternating small and large allocations") {
    for (int i = 0; i < 50; ++i) {
      void *small_ptr = arena.allocate_raw(16);
      void *large_ptr = arena.allocate_raw(256);
      CHECK(small_ptr != nullptr);
      CHECK(large_ptr != nullptr);
      CHECK(small_ptr != large_ptr);
    }
  }
}

TEST_CASE("ArenaAllocator move semantics") {
  ArenaAllocator arena1(512);
  arena1.allocate_raw(100);
  size_t used_bytes = arena1.used_bytes();
  size_t num_blocks = arena1.num_blocks();

  ArenaAllocator arena2 = std::move(arena1);
  CHECK(arena2.used_bytes() == used_bytes);
  CHECK(arena2.num_blocks() == num_blocks);

  void *ptr = arena2.allocate_raw(50);
  CHECK(ptr != nullptr);
}

TEST_CASE("ArenaAllocator edge cases") {
  SUBCASE("very small block size") {
    ArenaAllocator arena(16);
    void *ptr = arena.allocate_raw(8);
    CHECK(ptr != nullptr);
    CHECK(arena.num_blocks() == 1);
  }

  SUBCASE("allocation exactly block size") {
    ArenaAllocator arena(64);
    void *ptr = arena.allocate_raw(64);
    CHECK(ptr != nullptr);
    CHECK(arena.num_blocks() == 1);

    void *ptr2 = arena.allocate_raw(1);
    CHECK(ptr2 != nullptr);
    CHECK(arena.num_blocks() == 2);
  }

  SUBCASE("multiple resets") {
    ArenaAllocator arena;
    for (int i = 0; i < 10; ++i) {
      arena.allocate_raw(100);
      arena.reset();
      CHECK(arena.used_bytes() == 0);
    }
  }
}

struct TestPOD {
  int value;
  char name[16];

  TestPOD(int v, const char *n) : value(v) {
    std::strncpy(name, n, sizeof(name) - 1);
    name[sizeof(name) - 1] = '\0';
  }
};

TEST_CASE("ArenaAllocator with custom objects") {
  ArenaAllocator arena;

  TestPOD *obj1 = arena.construct<TestPOD>(42, "first");
  TestPOD *obj2 = arena.construct<TestPOD>(84, "second");

  CHECK(obj1 != nullptr);
  CHECK(obj2 != nullptr);
  CHECK(obj1 != obj2);
  CHECK(obj1->value == 42);
  CHECK(std::strcmp(obj1->name, "first") == 0);
  CHECK(obj2->value == 84);
  CHECK(std::strcmp(obj2->name, "second") == 0);
}

TEST_CASE("ArenaAllocator geometric growth policy") {
  ArenaAllocator arena(64);

  SUBCASE("normal geometric growth doubles size") {
    arena.allocate_raw(60); // Fill first block
    size_t initial_total = arena.total_allocated();

    arena.allocate_raw(10); // Force new block
    CHECK(arena.num_blocks() == 2);
    CHECK(arena.total_allocated() == initial_total + 128); // 64 * 2 = 128
  }

  SUBCASE("large allocation creates appropriately sized block") {
    arena.allocate_raw(60); // Fill first block
    size_t initial_total = arena.total_allocated();

    arena.allocate_raw(200); // Force large block
    CHECK(arena.num_blocks() == 2);
    CHECK(arena.total_allocated() >= initial_total + 200); // At least 200 bytes
  }

  SUBCASE("multiple growths maintain O(log n) blocks") {
    size_t allocation_size = 32;

    for (int i = 0; i < 10; ++i) {
      arena.allocate_raw(allocation_size);
    }

    // Should have grown logarithmically, not linearly
    CHECK(arena.num_blocks() < 6); // Much less than 10
  }
}

TEST_CASE("ArenaAllocator alignment edge cases") {
  ArenaAllocator arena;

  SUBCASE("unaligned then aligned allocation") {
    void *ptr1 = arena.allocate_raw(1, 1);
    void *ptr2 = arena.allocate_raw(8, 8);

    CHECK(ptr1 != nullptr);
    CHECK(ptr2 != nullptr);
    CHECK(reinterpret_cast<uintptr_t>(ptr2) % 8 == 0);
  }

  SUBCASE("large alignment requirements") {
    ArenaAllocator fresh_arena;
    void *ptr = fresh_arena.allocate_raw(1, 128);
    CHECK(ptr != nullptr);
    CHECK(reinterpret_cast<uintptr_t>(ptr) % 128 == 0);
  }
}

TEST_CASE("ArenaAllocator realloc functionality") {
  ArenaAllocator arena;

  SUBCASE("realloc edge cases") {
    // realloc with new_size == 0 returns nullptr and reclaims memory if it's
    // the last allocation
    ArenaAllocator fresh_arena(256);
    void *ptr = fresh_arena.allocate_raw(100);
    size_t used_before = fresh_arena.used_bytes();
    CHECK(used_before == 100);

    void *result = fresh_arena.realloc_raw(ptr, 100, 0);
    CHECK(result == nullptr);
    CHECK(fresh_arena.used_bytes() == 0); // Memory should be reclaimed

    // Test case where it's NOT the last allocation - memory cannot be reclaimed
    ArenaAllocator arena2(256);
    void *ptr1 = arena2.allocate_raw(50);
    (void)arena2.allocate_raw(50);
    size_t used_before2 = arena2.used_bytes();
    CHECK(used_before2 >= 100); // At least 100 bytes due to potential alignment

    void *result2 =
        arena2.realloc_raw(ptr1, 50, 0); // ptr1 is not the last allocation
    CHECK(result2 == nullptr);
    CHECK(arena2.used_bytes() ==
          used_before2); // Memory should NOT be reclaimed

    // realloc with nullptr behaves like allocate
    void *new_ptr = arena.realloc_raw(nullptr, 0, 150);
    CHECK(new_ptr != nullptr);
    CHECK(arena.used_bytes() >= 150);

    // realloc with same size returns same pointer
    void *same_ptr = arena.realloc_raw(new_ptr, 150, 150);
    CHECK(same_ptr == new_ptr);
  }

  SUBCASE("in-place extension - growing") {
    ArenaAllocator fresh_arena(1024);
    void *ptr = fresh_arena.allocate_raw(100);
    CHECK(ptr != nullptr);

    // Fill the allocation with test data
    std::memset(ptr, 0xAB, 100);

    // Should extend in place since it's the last allocation
    void *extended_ptr = fresh_arena.realloc_raw(ptr, 100, 200);
    CHECK(extended_ptr == ptr); // Should be same pointer (in-place)

    // Check that original data is preserved
    char *data = static_cast<char *>(extended_ptr);
    for (size_t i = 0; i < 100; ++i) {
      CHECK(data[i] == static_cast<char>(0xAB));
    }

    CHECK(fresh_arena.used_bytes() == 200);
  }

  SUBCASE("in-place shrinking") {
    ArenaAllocator fresh_arena(1024);
    void *ptr = fresh_arena.allocate_raw(200);
    std::memset(ptr, 0xCD, 200);

    // Should shrink in place
    void *shrunk_ptr = fresh_arena.realloc_raw(ptr, 200, 100);
    CHECK(shrunk_ptr == ptr); // Same pointer
    CHECK(fresh_arena.used_bytes() == 100);

    // Check that remaining data is preserved
    char *data = static_cast<char *>(shrunk_ptr);
    for (size_t i = 0; i < 100; ++i) {
      CHECK(data[i] == static_cast<char>(0xCD));
    }
  }

  SUBCASE("copy when can't extend in place") {
    ArenaAllocator fresh_arena(256); // Larger block to avoid edge cases

    // Allocate first chunk
    void *ptr1 = fresh_arena.allocate_raw(60);
    std::memset(ptr1, 0x11, 60);

    // Allocate second chunk (this prevents in-place extension of ptr1)
    void *ptr2 = fresh_arena.allocate_raw(30);
    std::memset(ptr2, 0x22, 30);

    // Try to reallocate ptr1 - should copy since ptr2 is now the last
    // allocation
    void *realloc_ptr1 = fresh_arena.realloc_raw(ptr1, 60, 120);
    CHECK(realloc_ptr1 != ptr1); // Should be different pointer (copy occurred)

    // Check that data was copied correctly
    char *data = static_cast<char *>(realloc_ptr1);
    for (size_t i = 0; i < 60; ++i) {
      CHECK(data[i] == static_cast<char>(0x11));
    }

    // Try to reallocate ptr2 (last allocation) - should extend in place if
    // space allows
    void *realloc_ptr2 = fresh_arena.realloc_raw(ptr2, 30, 50);
    // Note: this might not be in-place if the realloc of ptr1 used up space in
    // the block Let's just check the result is valid and data is preserved
    CHECK(realloc_ptr2 != nullptr);

    char *data2 = static_cast<char *>(realloc_ptr2);
    for (size_t i = 0; i < 30; ++i) {
      CHECK(data2[i] == static_cast<char>(0x22));
    }
  }

  SUBCASE("copy when insufficient space for extension") {
    ArenaAllocator fresh_arena(100);

    // Allocate almost all space
    void *ptr = fresh_arena.allocate_raw(90);
    std::memset(ptr, 0x33, 90);

    // Try to extend beyond block size - should copy to new block
    void *extended_ptr = fresh_arena.realloc_raw(ptr, 90, 150);
    CHECK(extended_ptr != ptr);           // Should be different (new block)
    CHECK(fresh_arena.num_blocks() == 2); // Should have created new block

    // Check data preservation
    char *data = static_cast<char *>(extended_ptr);
    for (size_t i = 0; i < 90; ++i) {
      CHECK(data[i] == static_cast<char>(0x33));
    }
  }

  SUBCASE("realloc with custom alignment") {
    ArenaAllocator fresh_arena(1024);

    // Allocate with specific alignment
    void *ptr = fresh_arena.allocate_raw(50, 16);
    CHECK(reinterpret_cast<uintptr_t>(ptr) % 16 == 0);
    std::memset(ptr, 0x44, 50);

    // Realloc with same alignment - should extend in place
    void *extended_ptr = fresh_arena.realloc_raw(ptr, 50, 100, 16);
    CHECK(extended_ptr == ptr); // In place
    CHECK(reinterpret_cast<uintptr_t>(extended_ptr) % 16 == 0);

    // Check data preservation
    char *data = static_cast<char *>(extended_ptr);
    for (size_t i = 0; i < 50; ++i) {
      CHECK(data[i] == static_cast<char>(0x44));
    }
  }

  SUBCASE("realloc stress test") {
    ArenaAllocator fresh_arena(512);
    void *ptr = fresh_arena.allocate_raw(50);
    size_t current_size = 50;

    // Fill with pattern
    for (size_t i = 0; i < 50; ++i) {
      static_cast<char *>(ptr)[i] = static_cast<char>(i & 0xFF);
    }

    // Grow in multiple steps
    for (size_t new_size = 100; new_size <= 500; new_size += 100) {
      void *new_ptr = fresh_arena.realloc_raw(ptr, current_size, new_size);
      CHECK(new_ptr != nullptr);

      // Verify original data is still there
      for (size_t i = 0; i < 50; ++i) {
        CHECK(static_cast<char *>(new_ptr)[i] == static_cast<char>(i & 0xFF));
      }

      ptr = new_ptr;
      current_size = new_size;
    }
  }
}
