#include <memory>
#include <thread>
#include <vector>

#include <doctest/doctest.h>
#include <nanobench.h>

#include "reference.hpp"

namespace {

struct TestObject {
  int64_t data = 42;

  TestObject() = default;
  explicit TestObject(int64_t value) : data(value) {}
};

// Trait helpers for templated benchmarks
template <typename T> struct PointerTraits;

template <typename T> struct PointerTraits<std::shared_ptr<T>> {
  using pointer_type = std::shared_ptr<T>;
  using weak_type = std::weak_ptr<T>;

  template <typename... Args> static pointer_type make(Args &&...args) {
    return std::make_shared<T>(std::forward<Args>(args)...);
  }

  static const char *name() { return "std::shared_ptr"; }
  static const char *weak_name() { return "std::weak_ptr"; }
};

template <typename T> struct PointerTraits<Ref<T>> {
  using pointer_type = Ref<T>;
  using weak_type = WeakRef<T>;

  template <typename... Args> static pointer_type make(Args &&...args) {
    return make_ref<T>(std::forward<Args>(args)...);
  }

  static const char *name() { return "Ref"; }
  static const char *weak_name() { return "WeakRef"; }
};

// Force multi-threaded mode to defeat __libc_single_threaded optimization
void force_multithreaded() {
  std::thread t([]() {});
  t.join();
}

template <typename PtrType>
void benchmark_creation(ankerl::nanobench::Bench &bench) {
  using Traits = PointerTraits<PtrType>;
  force_multithreaded();

  bench.run(std::string(Traits::name()) + " creation", [&] {
    auto ptr = Traits::make(TestObject{123});
    ankerl::nanobench::doNotOptimizeAway(ptr);
  });
}

template <typename PtrType>
void benchmark_copy(ankerl::nanobench::Bench &bench) {
  using Traits = PointerTraits<PtrType>;
  force_multithreaded();

  auto original = Traits::make(TestObject{123});
  bench.run(std::string(Traits::name()) + " copy", [&] {
    auto copy = original;
    ankerl::nanobench::doNotOptimizeAway(copy);
  });
}

template <typename PtrType>
void benchmark_move(ankerl::nanobench::Bench &bench) {
  using Traits = PointerTraits<PtrType>;

  auto original = Traits::make(TestObject{123});
  bench.run(std::string(Traits::name()) + " move", [&] {
    auto moved = std::move(original);
    ankerl::nanobench::doNotOptimizeAway(moved);
    original = std::move(moved);
    ankerl::nanobench::doNotOptimizeAway(original);
  });
}

template <typename PtrType>
void benchmark_weak_copy(ankerl::nanobench::Bench &bench) {
  using Traits = PointerTraits<PtrType>;
  force_multithreaded();

  auto strong_ptr = Traits::make(TestObject{123});
  typename Traits::weak_type weak_original = strong_ptr;
  bench.run(std::string(Traits::weak_name()) + " copy", [&] {
    auto weak_copy = weak_original;
    ankerl::nanobench::doNotOptimizeAway(weak_copy);
  });
}

template <typename PtrType>
void benchmark_weak_move(ankerl::nanobench::Bench &bench) {
  using Traits = PointerTraits<PtrType>;

  auto strong_ptr = Traits::make(TestObject{123});
  typename Traits::weak_type weak_original = strong_ptr;
  bench.run(std::string(Traits::weak_name()) + " move", [&] {
    auto weak_moved = std::move(weak_original);
    ankerl::nanobench::doNotOptimizeAway(weak_moved);
    weak_original = std::move(weak_moved);
    ankerl::nanobench::doNotOptimizeAway(weak_original);
  });
}

template <typename PtrType>
void benchmark_dereference(ankerl::nanobench::Bench &bench) {
  using Traits = PointerTraits<PtrType>;

  auto ptr = Traits::make(TestObject{456});
  bench.run(std::string(Traits::name()) + " dereference",
            [&] { ankerl::nanobench::doNotOptimizeAway(ptr->data); });
}

template <typename PtrType>
void benchmark_weak_lock_success(ankerl::nanobench::Bench &bench) {
  using Traits = PointerTraits<PtrType>;

  auto strong_ptr = Traits::make(TestObject{789});
  typename Traits::weak_type weak_ptr = strong_ptr;
  bench.run(std::string(Traits::weak_name()) + " lock success", [&] {
    auto locked = weak_ptr.lock();
    ankerl::nanobench::doNotOptimizeAway(locked);
  });
}

template <typename PtrType>
void benchmark_weak_lock_failure(ankerl::nanobench::Bench &bench) {
  using Traits = PointerTraits<PtrType>;

  typename Traits::weak_type weak_ptr;
  {
    auto strong_ptr = Traits::make(TestObject{999});
    weak_ptr = strong_ptr;
  }
  bench.run(std::string(Traits::weak_name()) + " lock failure", [&] {
    auto locked = weak_ptr.lock();
    ankerl::nanobench::doNotOptimizeAway(locked);
  });
}

template <typename PtrType>
void benchmark_multithreaded_copy(ankerl::nanobench::Bench &bench,
                                  int num_threads) {
  using Traits = PointerTraits<PtrType>;

  // Create the shared object outside the benchmark
  auto ptr = Traits::make(TestObject{456});

  // Create background threads that will create contention
  std::atomic<bool> keep_running{true};
  std::vector<std::thread> background_threads;

  for (int i = 0; i < num_threads - 1; ++i) {
    background_threads.emplace_back([&]() {
      while (keep_running.load(std::memory_order_relaxed)) {
        auto copy = ptr;
        ankerl::nanobench::doNotOptimizeAway(copy);
      }
    });
  }

  // Benchmark the foreground thread under contention
  bench.run(std::string(Traits::name()) + " copy under contention", [&] {
    auto copy = ptr;
    ankerl::nanobench::doNotOptimizeAway(copy);
  });

  // Clean up background threads
  keep_running.store(false, std::memory_order_relaxed);
  for (auto &t : background_threads) {
    t.join();
  }
}

template <typename PtrType>
void benchmark_multithreaded_weak_lock(ankerl::nanobench::Bench &bench,
                                       int num_threads) {
  using Traits = PointerTraits<PtrType>;

  // Create the shared object and weak reference outside the benchmark
  auto strong_ptr = Traits::make(TestObject{789});
  typename Traits::weak_type weak_ptr = strong_ptr;

  // Create background threads that will create contention
  std::atomic<bool> keep_running{true};
  std::vector<std::thread> background_threads;

  for (int i = 0; i < num_threads - 1; ++i) {
    background_threads.emplace_back([&]() {
      while (keep_running.load(std::memory_order_relaxed)) {
        auto locked = weak_ptr.lock();
        ankerl::nanobench::doNotOptimizeAway(locked);
      }
    });
  }

  // Benchmark the foreground thread under contention
  bench.run(std::string(Traits::weak_name()) + " lock under contention", [&] {
    auto locked = weak_ptr.lock();
    ankerl::nanobench::doNotOptimizeAway(locked);
  });

  // Clean up background threads
  keep_running.store(false, std::memory_order_relaxed);
  for (auto &t : background_threads) {
    t.join();
  }
}

template <typename PtrType>
void benchmark_weak_copy_with_strong_contention(ankerl::nanobench::Bench &bench,
                                                int num_threads) {
  using Traits = PointerTraits<PtrType>;

  // Create the shared object and weak reference outside the benchmark
  auto strong_ptr = Traits::make(TestObject{456});
  typename Traits::weak_type weak_ptr = strong_ptr;

  // Create background threads copying the strong pointer
  std::atomic<bool> keep_running{true};
  std::vector<std::thread> background_threads;

  for (int i = 0; i < num_threads - 1; ++i) {
    background_threads.emplace_back([&]() {
      while (keep_running.load(std::memory_order_relaxed)) {
        auto copy = strong_ptr;
        ankerl::nanobench::doNotOptimizeAway(copy);
      }
    });
  }

  // Benchmark weak reference copying under strong reference contention
  bench.run(std::string(Traits::weak_name()) + " copy with strong contention",
            [&] {
              auto weak_copy = weak_ptr;
              ankerl::nanobench::doNotOptimizeAway(weak_copy);
            });

  // Clean up background threads
  keep_running.store(false, std::memory_order_relaxed);
  for (auto &t : background_threads) {
    t.join();
  }
}

template <typename PtrType>
void benchmark_strong_copy_with_weak_contention(ankerl::nanobench::Bench &bench,
                                                int num_threads) {
  using Traits = PointerTraits<PtrType>;

  // Create the shared object and weak reference outside the benchmark
  auto strong_ptr = Traits::make(TestObject{789});
  typename Traits::weak_type weak_ptr = strong_ptr;

  // Create background threads copying the weak pointer
  std::atomic<bool> keep_running{true};
  std::vector<std::thread> background_threads;

  for (int i = 0; i < num_threads - 1; ++i) {
    background_threads.emplace_back([&]() {
      while (keep_running.load(std::memory_order_relaxed)) {
        auto weak_copy = weak_ptr;
        ankerl::nanobench::doNotOptimizeAway(weak_copy);
      }
    });
  }

  // Benchmark strong reference copying under weak reference contention
  bench.run(std::string(Traits::name()) + " copy with weak contention", [&] {
    auto strong_copy = strong_ptr;
    ankerl::nanobench::doNotOptimizeAway(strong_copy);
  });

  // Clean up background threads
  keep_running.store(false, std::memory_order_relaxed);
  for (auto &t : background_threads) {
    t.join();
  }
}

} // anonymous namespace

TEST_CASE("Creation performance comparison") {
  ankerl::nanobench::Bench bench;
  bench.title("Creation performance comparison");
  bench.relative(true);
  benchmark_creation<std::shared_ptr<TestObject>>(bench);
  benchmark_creation<Ref<TestObject>>(bench);
}

TEST_CASE("Copy performance comparison") {
  ankerl::nanobench::Bench bench;
  bench.title("Copy performance comparison");
  bench.relative(true);
  benchmark_copy<std::shared_ptr<TestObject>>(bench);
  benchmark_copy<Ref<TestObject>>(bench);
}

TEST_CASE("Move performance comparison") {
  ankerl::nanobench::Bench bench;
  bench.title("Move performance comparison");
  bench.relative(true);
  benchmark_move<std::shared_ptr<TestObject>>(bench);
  benchmark_move<Ref<TestObject>>(bench);
}

TEST_CASE("Weak copy performance comparison") {
  ankerl::nanobench::Bench bench;
  bench.title("Weak copy performance comparison");
  bench.relative(true);
  benchmark_weak_copy<std::shared_ptr<TestObject>>(bench);
  benchmark_weak_copy<Ref<TestObject>>(bench);
}

TEST_CASE("Weak move performance comparison") {
  ankerl::nanobench::Bench bench;
  bench.title("Weak move performance comparison");
  bench.relative(true);
  benchmark_weak_move<std::shared_ptr<TestObject>>(bench);
  benchmark_weak_move<Ref<TestObject>>(bench);
}

TEST_CASE("Dereference performance comparison") {
  ankerl::nanobench::Bench bench;
  bench.title("Dereference performance comparison");
  bench.relative(true);
  benchmark_dereference<std::shared_ptr<TestObject>>(bench);
  benchmark_dereference<Ref<TestObject>>(bench);
}

TEST_CASE("Weak lock success performance comparison") {
  ankerl::nanobench::Bench bench;
  bench.title("Weak lock success performance comparison");
  bench.relative(true);
  benchmark_weak_lock_success<std::shared_ptr<TestObject>>(bench);
  benchmark_weak_lock_success<Ref<TestObject>>(bench);
}

TEST_CASE("Weak lock failure performance comparison") {
  ankerl::nanobench::Bench bench;
  bench.title("Weak lock failure performance comparison");
  bench.relative(true);
  benchmark_weak_lock_failure<std::shared_ptr<TestObject>>(bench);
  benchmark_weak_lock_failure<Ref<TestObject>>(bench);
}

TEST_CASE("Copy performance under contention") {
  const int num_threads = 3;
  ankerl::nanobench::Bench bench;
  bench.title("Copy performance under contention");
  bench.relative(true);
  bench.minEpochIterations(500000);
  benchmark_multithreaded_copy<std::shared_ptr<TestObject>>(bench, num_threads);
  benchmark_multithreaded_copy<Ref<TestObject>>(bench, num_threads);
}

TEST_CASE("Weak lock performance under contention") {
  const int num_threads = 3;
  ankerl::nanobench::Bench bench;
  bench.title("Weak lock performance under contention");
  bench.relative(true);
  bench.minEpochIterations(500000);
  benchmark_multithreaded_weak_lock<std::shared_ptr<TestObject>>(bench,
                                                                 num_threads);
  benchmark_multithreaded_weak_lock<Ref<TestObject>>(bench, num_threads);
}

TEST_CASE("Weak copy performance under strong reference contention") {
  const int num_threads = 3;
  ankerl::nanobench::Bench bench;
  bench.title("Weak copy performance under strong reference contention");
  bench.relative(true);
  bench.minEpochIterations(500000);
  benchmark_weak_copy_with_strong_contention<std::shared_ptr<TestObject>>(
      bench, num_threads);
  benchmark_weak_copy_with_strong_contention<Ref<TestObject>>(bench,
                                                              num_threads);
}

TEST_CASE("Strong copy performance under weak reference contention") {
  const int num_threads = 3;
  ankerl::nanobench::Bench bench;
  bench.title("Strong copy performance under weak reference contention");
  bench.relative(true);
  bench.minEpochIterations(500000);
  benchmark_strong_copy_with_weak_contention<std::shared_ptr<TestObject>>(
      bench, num_threads);
  benchmark_strong_copy_with_weak_contention<Ref<TestObject>>(bench,
                                                              num_threads);
}