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

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

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

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

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

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

  SUBCASE("allocate multiple bytes") {
    void *ptr1 = arena.allocate(100);
    void *ptr2 = arena.allocate(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(1);
    CHECK(reinterpret_cast<uintptr_t>(ptr) % alignof(std::max_align_t) == 0);
  }

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

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

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

  SUBCASE("alignment with larger allocations") {
    ArenaAllocator fresh_arena;
    void *ptr = fresh_arena.allocate(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(64);
    CHECK(ptr != nullptr);
    CHECK(arena.num_blocks() == 1);
    CHECK(arena.used_bytes() == 64);
  }

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

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

  SUBCASE("allocation larger than block size throws") {
    CHECK_THROWS_AS(arena.allocate(200), std::bad_alloc);
  }
}

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

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

  SUBCASE("construct string") {
    std::string *ptr = arena.construct<std::string>("hello world");
    CHECK(ptr != nullptr);
    CHECK(*ptr == "hello world");
  }

  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") {
    auto *ptr = arena.construct<std::pair<int, std::string>>(42, "test");
    CHECK(ptr != nullptr);
    CHECK(ptr->first == 42);
    CHECK(ptr->second == "test");
  }
}

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

  arena.allocate(100);
  arena.allocate(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(50);
  CHECK(ptr != nullptr);
  CHECK(arena.used_bytes() == 50);
}

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

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

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

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

  arena.allocate(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(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(16);
      void *large_ptr = arena.allocate(256);
      CHECK(small_ptr != nullptr);
      CHECK(large_ptr != nullptr);
      CHECK(small_ptr != large_ptr);
    }
  }
}

TEST_CASE("ArenaAllocator move semantics") {
  ArenaAllocator arena1(512);
  arena1.allocate(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(50);
  CHECK(ptr != nullptr);
}

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

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

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

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

struct TestObject {
  int value;
  std::string name;

  TestObject(int v, const std::string &n) : value(v), name(n) {}
  ~TestObject() = default;
};

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

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

  CHECK(obj1 != nullptr);
  CHECK(obj2 != nullptr);
  CHECK(obj1 != obj2);
  CHECK(obj1->value == 42);
  CHECK(obj1->name == "first");
  CHECK(obj2->value == 84);
  CHECK(obj2->name == "second");
}

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

  SUBCASE("unaligned then aligned allocation") {
    void *ptr1 = arena.allocate(1, 1);
    void *ptr2 = arena.allocate(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(1, 128);
    CHECK(ptr != nullptr);
    CHECK(reinterpret_cast<uintptr_t>(ptr) % 128 == 0);
  }
}