#include <Utils/Containers/SlotMap.hpp>

#include <ankerl/unordered_dense.h>
#include <benchmark/benchmark.h>

#include <algorithm>
#include <cstdint>
#include <random>
#include <type_traits>
#include <utility>
#include <vector>

namespace Bigfoot
{
struct MyComplexStruct
{
    std::uint32_t m_actualValueForBench = 0;   // field summed by the Iterate benchmark

    float m_position[3] = {0.0f, 0.0f, 0.0f};
    float m_rotation[4] = {0.0f, 0.0f, 0.0f, 1.0f};
    float m_scale[3] = {1.0f, 1.0f, 1.0f};

    std::uint64_t m_entityId = 0;
    std::uint32_t m_parentSlot = 0;
    std::uint32_t m_flags = 0;

    MyComplexStruct() = default;

    explicit MyComplexStruct(const std::uint32_t p_value):
        m_actualValueForBench(p_value),
        m_entityId(p_value)
    {
    }
};

static_assert(sizeof(MyComplexStruct) == 64, "Payload should be exactly one cache line.");
static_assert(std::is_trivially_copyable_v<MyComplexStruct>,
              "Container moves should be a plain memcpy, not a heap-touching move.");

class SlotMapAdaptor
{
  public:
    using Key = SlotMap<MyComplexStruct>::SlotKey;

    Key Add(MyComplexStruct&& p_value)
    {
        return m_container.Insert(std::move(p_value));
    }

    const MyComplexStruct* Find(const Key p_key) const
    {
        return m_container.Get(p_key);
    }

    void Remove(const Key p_key)
    {
        m_container.Remove(p_key);
    }

    void Clear()
    {
        m_container.Reset();
    }

    template<class FUNC>
    void ForEachValue(FUNC&& p_func) const
    {
        for (const MyComplexStruct& value: m_container)
        {
            p_func(value);
        }
    }

  private:
    SlotMap<MyComplexStruct> m_container;
};

template<class MAP>
class HashMapHarnessAdaptor
{
  public:
    using Key = typename MAP::key_type;

    Key Add(MyComplexStruct&& p_value)
    {
        const Key key = m_nextKey++;
        m_container.emplace(key, std::move(p_value));
        return key;
    }

    const MyComplexStruct* Find(const Key p_key) const
    {
        const auto it = m_container.find(p_key);
        return it != m_container.end() ? &it->second : nullptr;
    }

    void Remove(const Key p_key)
    {
        m_container.erase(p_key);
    }

    void Clear()
    {
        m_container.clear();
        m_nextKey = 0;
    }

    template<class FUNC>
    void ForEachValue(FUNC&& p_func) const
    {
        for (const auto& keyValue: m_container)
        {
            p_func(keyValue.second);
        }
    }

  private:
    MAP m_container;
    Key m_nextKey = 0;
};

template<class ADAPTOR>
void Insert(benchmark::State& state)
{
    const auto count = static_cast<std::uint32_t>(state.range(0));

    for (auto _: state)
    {
        ADAPTOR adaptor;
        for (std::uint32_t i = 0; i < count; ++i)
        {
            benchmark::DoNotOptimize(adaptor.Add(MyComplexStruct {i}));
        }
        benchmark::DoNotOptimize(adaptor);
        benchmark::ClobberMemory();
    }

    state.SetItemsProcessed(state.iterations() * count);
}

template<class ADAPTOR>
void Remove(benchmark::State& state)
{
    const auto count = static_cast<std::uint32_t>(state.range(0));

    ADAPTOR adaptor;
    std::vector<typename ADAPTOR::Key> keys;
    keys.reserve(count);
    std::mt19937 rng(0x5EEDU);

    for (auto _: state)
    {
        state.PauseTiming();
        adaptor.Clear();
        keys.clear();
        for (std::uint32_t i = 0; i < count; ++i)
        {
            keys.push_back(adaptor.Add(MyComplexStruct {i}));
        }
        std::shuffle(keys.begin(), keys.end(), rng);
        state.ResumeTiming();

        for (const typename ADAPTOR::Key key: keys)
        {
            adaptor.Remove(key);
        }
        benchmark::DoNotOptimize(adaptor);
        benchmark::ClobberMemory();
    }

    state.SetItemsProcessed(state.iterations() * count);
}

template<class ADAPTOR>
void Get(benchmark::State& state)
{
    const auto count = static_cast<std::uint32_t>(state.range(0));

    ADAPTOR adaptor;
    std::vector<typename ADAPTOR::Key> keys;
    keys.reserve(count);
    for (std::uint32_t i = 0; i < count; ++i)
    {
        keys.push_back(adaptor.Add(MyComplexStruct {i}));
    }
    std::mt19937 rng(0x5EEDU);
    std::shuffle(keys.begin(), keys.end(), rng);

    for (auto _: state)
    {
        std::uint64_t sum = 0;
        for (const typename ADAPTOR::Key key: keys)
        {
            if (const MyComplexStruct* const value = adaptor.Find(key))
            {
                sum += value->m_actualValueForBench;
            }
        }
        benchmark::DoNotOptimize(sum);
    }

    state.SetItemsProcessed(state.iterations() * count);
}

template<class ADAPTOR>
void Iterate(benchmark::State& state)
{
    const auto count = static_cast<std::uint32_t>(state.range(0));

    ADAPTOR adaptor;
    for (std::uint32_t i = 0; i < count; ++i)
    {
        adaptor.Add(MyComplexStruct {i});
    }

    for (auto _: state)
    {
        std::uint64_t sum = 0;
        adaptor.ForEachValue(
            [&sum](const MyComplexStruct& value)
            {
                sum += value.m_actualValueForBench;
            });
        benchmark::DoNotOptimize(sum);
    }

    state.SetItemsProcessed(state.iterations() * count);
}

// Register all four benchmarks for one container adapter under a readable name.
#define BIGFOOT_REGISTER_BENCHMARKS(ADAPTOR, NAME)                                                                     \
    BENCHMARK_TEMPLATE(Insert, ADAPTOR)->Name(NAME "/Insert")->Range(8, 8 << 20);                                      \
    BENCHMARK_TEMPLATE(Remove, ADAPTOR)->Name(NAME "/Remove")->Range(8, 8 << 20);                                      \
    BENCHMARK_TEMPLATE(Get, ADAPTOR)->Name(NAME "/Get")->Range(8, 8 << 20);                                            \
    BENCHMARK_TEMPLATE(Iterate, ADAPTOR)->Name(NAME "/Iterate")->Range(8, 8 << 20)

using UnorderedDenseMapHarness = HashMapHarnessAdaptor<ankerl::unordered_dense::map<std::uint64_t, MyComplexStruct>>;
using UnorderedDenseSegementedMapHarness =
    HashMapHarnessAdaptor<ankerl::unordered_dense::segmented_map<std::uint64_t, MyComplexStruct>>;
using UnorderedMapHarness = HashMapHarnessAdaptor<std::unordered_map<std::uint64_t, MyComplexStruct>>;
using MapHarness = HashMapHarnessAdaptor<std::map<std::uint64_t, MyComplexStruct>>;

BIGFOOT_REGISTER_BENCHMARKS(SlotMapAdaptor, "Bigfoot::SlotMap");
BIGFOOT_REGISTER_BENCHMARKS(UnorderedDenseMapHarness, "ankerl::unordered_dense::map");
BIGFOOT_REGISTER_BENCHMARKS(UnorderedDenseSegementedMapHarness, "ankerl::unordered_dense::segmented_map");
BIGFOOT_REGISTER_BENCHMARKS(UnorderedMapHarness, "std::unordered_map");
BIGFOOT_REGISTER_BENCHMARKS(MapHarness, "std::map");

#undef BIGFOOT_REGISTER_BENCHMARKS
}   // namespace Bigfoot