#include "arena_allocator.hpp"
#include <algorithm>
#include <iostream>
#include <memory>
#include <nanobench.h>
#include <string>
#include <vector>

struct TestStruct {
  int a;
  double b;
  char c[64];

  TestStruct(int x, double y) : a(x), b(y) {
    std::fill(std::begin(c), std::end(c), 'x');
  }
};

class BenchmarkResults {
public:
  void add_result(const std::string &name, double ops_per_sec) {
    results_.push_back({name, ops_per_sec});
  }

  void print_summary() {
    std::cout << "\n=== Arena Allocator Benchmark Summary ===\n";
    for (const auto &result : results_) {
      std::cout << result.first << ": " << result.second << " ops/sec\n";
    }
    std::cout << "==========================================\n";
  }

private:
  std::vector<std::pair<std::string, double>> results_;
};

int main() {
  BenchmarkResults results;

  // Small allocation benchmark - Arena vs malloc
  {
    constexpr size_t NUM_ALLOCS = 10000;
    constexpr size_t ALLOC_SIZE = 32;

    auto bench = ankerl::nanobench::Bench()
                     .title("Small Allocations (32 bytes)")
                     .unit("allocation")
                     .warmup(100)
                     .epochs(1000);

    // Arena allocator benchmark
    bench.run("ArenaAllocator", [&] {
      ArenaAllocator arena(4 * 1024 * 1024);
      for (size_t i = 0; i < NUM_ALLOCS; ++i) {
        void *ptr = arena.allocate(ALLOC_SIZE);
        ankerl::nanobench::doNotOptimizeAway(ptr);
      }
    });

    // Standard malloc benchmark
    std::vector<void *> malloc_ptrs;
    malloc_ptrs.reserve(NUM_ALLOCS);

    bench.run("malloc", [&] {
      malloc_ptrs.clear();
      for (size_t i = 0; i < NUM_ALLOCS; ++i) {
        void *ptr = std::malloc(ALLOC_SIZE);
        malloc_ptrs.push_back(ptr);
        ankerl::nanobench::doNotOptimizeAway(ptr);
      }

      for (void *ptr : malloc_ptrs) {
        std::free(ptr);
      }
    });
  }

  // Medium allocation benchmark
  {
    constexpr size_t NUM_ALLOCS = 1000;
    constexpr size_t ALLOC_SIZE = 1024;

    auto bench = ankerl::nanobench::Bench()
                     .title("Medium Allocations (1024 bytes)")
                     .unit("allocation")
                     .warmup(50)
                     .epochs(500);

    bench.run("ArenaAllocator", [&] {
      ArenaAllocator arena(4 * 1024 * 1024);
      for (size_t i = 0; i < NUM_ALLOCS; ++i) {
        void *ptr = arena.allocate(ALLOC_SIZE);
        ankerl::nanobench::doNotOptimizeAway(ptr);
      }
    });

    std::vector<void *> malloc_ptrs;
    malloc_ptrs.reserve(NUM_ALLOCS);

    bench.run("malloc", [&] {
      malloc_ptrs.clear();
      for (size_t i = 0; i < NUM_ALLOCS; ++i) {
        void *ptr = std::malloc(ALLOC_SIZE);
        malloc_ptrs.push_back(ptr);
        ankerl::nanobench::doNotOptimizeAway(ptr);
      }

      for (void *ptr : malloc_ptrs) {
        std::free(ptr);
      }
    });
  }

  // Object construction benchmark
  {
    constexpr size_t NUM_CONSTRUCTS = 5000;

    auto bench = ankerl::nanobench::Bench()
                     .title("Object Construction")
                     .unit("construction")
                     .warmup(50)
                     .epochs(500);

    bench.run("ArenaAllocator::construct", [&] {
      ArenaAllocator arena(4 * 1024 * 1024);
      for (size_t i = 0; i < NUM_CONSTRUCTS; ++i) {
        TestStruct *obj = arena.construct<TestStruct>(i, i * 1.5);
        ankerl::nanobench::doNotOptimizeAway(obj);
      }
    });

    std::vector<std::unique_ptr<TestStruct>> objects;
    objects.reserve(NUM_CONSTRUCTS);

    bench.run("std::make_unique", [&] {
      objects.clear();
      for (size_t i = 0; i < NUM_CONSTRUCTS; ++i) {
        auto obj = std::make_unique<TestStruct>(i, i * 1.5);
        ankerl::nanobench::doNotOptimizeAway(obj.get());
        objects.push_back(std::move(obj));
      }
    });
  }

  // String allocation benchmark
  {
    constexpr size_t NUM_STRINGS = 1000;
    const std::vector<std::string> test_strings = {
        "short", "medium length string for testing",
        "this is a much longer string that should test the allocation "
        "performance with larger objects and see how well the arena allocator "
        "handles variable sized allocations"};

    auto bench = ankerl::nanobench::Bench()
                     .title("String Construction")
                     .unit("string")
                     .warmup(50)
                     .epochs(300);

    bench.run("ArenaAllocator", [&] {
      ArenaAllocator arena(4 * 1024 * 1024);
      for (size_t i = 0; i < NUM_STRINGS; ++i) {
        const auto &test_str = test_strings[i % test_strings.size()];
        std::string *str = arena.construct<std::string>(test_str);
        ankerl::nanobench::doNotOptimizeAway(str);
      }
    });

    std::vector<std::unique_ptr<std::string>> strings;
    strings.reserve(NUM_STRINGS);

    bench.run("std::make_unique<string>", [&] {
      strings.clear();
      for (size_t i = 0; i < NUM_STRINGS; ++i) {
        const auto &test_str = test_strings[i % test_strings.size()];
        auto str = std::make_unique<std::string>(test_str);
        ankerl::nanobench::doNotOptimizeAway(str.get());
        strings.push_back(std::move(str));
      }
    });
  }

  // Mixed size allocation pattern
  {
    constexpr size_t NUM_ALLOCS = 2000;
    const std::vector<size_t> sizes = {8, 16, 32, 64, 128, 256, 512, 1024};

    auto bench = ankerl::nanobench::Bench()
                     .title("Mixed Size Allocations")
                     .unit("allocation")
                     .warmup(50)
                     .epochs(300);

    bench.run("ArenaAllocator", [&] {
      ArenaAllocator arena(4 * 1024 * 1024);
      for (size_t i = 0; i < NUM_ALLOCS; ++i) {
        size_t size = sizes[i % sizes.size()];
        void *ptr = arena.allocate(size);
        ankerl::nanobench::doNotOptimizeAway(ptr);
      }
    });

    std::vector<void *> malloc_ptrs;
    malloc_ptrs.reserve(NUM_ALLOCS);

    bench.run("malloc", [&] {
      malloc_ptrs.clear();
      for (size_t i = 0; i < NUM_ALLOCS; ++i) {
        size_t size = sizes[i % sizes.size()];
        void *ptr = std::malloc(size);
        malloc_ptrs.push_back(ptr);
        ankerl::nanobench::doNotOptimizeAway(ptr);
      }

      for (void *ptr : malloc_ptrs) {
        std::free(ptr);
      }
    });
  }

  // Arena reset performance
  {
    constexpr size_t NUM_RESETS = 1000;
    constexpr size_t ALLOCS_PER_RESET = 100;

    auto bench = ankerl::nanobench::Bench()
                     .title("Arena Reset Performance")
                     .unit("reset")
                     .warmup(20)
                     .epochs(200);

    bench.run("ArenaAllocator reset", [&] {
      ArenaAllocator arena(64 * 1024);
      for (size_t i = 0; i < NUM_RESETS; ++i) {
        // Allocate some memory
        for (size_t j = 0; j < ALLOCS_PER_RESET; ++j) {
          void *ptr = arena.allocate(64);
          ankerl::nanobench::doNotOptimizeAway(ptr);
        }
        // Reset the arena
        arena.reset();
        ankerl::nanobench::doNotOptimizeAway(&arena);
      }
    });
  }

  // Alignment performance test
  {
    constexpr size_t NUM_ALLOCS = 5000;
    const std::vector<size_t> alignments = {8, 16, 32, 64, 128};

    auto bench = ankerl::nanobench::Bench()
                     .title("Aligned Allocations")
                     .unit("allocation")
                     .warmup(50)
                     .epochs(300);

    bench.run("ArenaAllocator aligned", [&] {
      ArenaAllocator arena(4 * 1024 * 1024);
      for (size_t i = 0; i < NUM_ALLOCS; ++i) {
        size_t alignment = alignments[i % alignments.size()];
        void *ptr = arena.allocate(64, alignment);
        ankerl::nanobench::doNotOptimizeAway(ptr);
      }
    });

    std::vector<void *> aligned_ptrs;
    aligned_ptrs.reserve(NUM_ALLOCS);

    bench.run("aligned_alloc", [&] {
      aligned_ptrs.clear();
      for (size_t i = 0; i < NUM_ALLOCS; ++i) {
        size_t alignment = alignments[i % alignments.size()];
        void *ptr = std::aligned_alloc(alignment, 64);
        aligned_ptrs.push_back(ptr);
        ankerl::nanobench::doNotOptimizeAway(ptr);
      }

      for (void *ptr : aligned_ptrs) {
        std::free(ptr);
      }
    });
  }

  // Block growth performance
  {
    constexpr size_t INITIAL_BLOCK_SIZE = 1024;
    constexpr size_t NUM_LARGE_ALLOCS = 10;
    constexpr size_t LARGE_ALLOC_SIZE = 512;

    auto bench = ankerl::nanobench::Bench()
                     .title("Block Growth Performance")
                     .unit("allocation")
                     .warmup(20)
                     .epochs(100);

    bench.run("ArenaAllocator block growth", [&] {
      ArenaAllocator arena(INITIAL_BLOCK_SIZE);
      for (size_t i = 0; i < NUM_LARGE_ALLOCS; ++i) {
        void *ptr = arena.allocate(LARGE_ALLOC_SIZE);
        ankerl::nanobench::doNotOptimizeAway(ptr);
      }
    });
  }

  std::cout << "\nBenchmarks completed successfully!\n";

  return 0;
}