#include <kapi/devices.hpp>

#include <kstd/format.hpp>
#include <kstd/memory.hpp>

#include <catch2/catch_test_macros.hpp>

#include <atomic>
#include <cstddef>
#include <mutex>
#include <thread>
#include <tuple>
#include <vector>

namespace
{

  struct test_device final : kapi::devices::device
  {
    using kapi::devices::device::device;
  };

}  // namespace

constexpr auto thread_count = 8uz;
constexpr auto devices_per_thread = 200uz;

SCENARIO("Concurrent attach/detach/lookup on a bus is race-free")
{
  GIVEN("A bus shared by several threads")
  {
    auto shared_bus = kstd::make_shared<kapi::devices::bus>("stress_test_bus");
    kapi::devices::get_root_bus()->add_child(shared_bus);

    WHEN("each thread repeatedly attaches, detaches, and looks up devices concurrently with the others")
    {
      auto threads = std::vector<std::jthread>{};
      threads.reserve(thread_count);

      auto failure_count = std::atomic<std::size_t>{0};

      for (auto thread_index = 0uz; thread_index < thread_count; ++thread_index)
      {
        threads.emplace_back([&shared_bus, &failure_count, thread_index] {
          for (auto i = 0uz; i < devices_per_thread; ++i)
          {
            auto name = kstd::format("stress_test_device_{}_{}", thread_index, i);
            auto device = kstd::make_shared<test_device>(name);

            shared_bus->add_child(device);

            std::ignore = kapi::devices::device_registry::get().find(name);
            std::ignore = kapi::devices::device_registry::get().all();
            std::ignore = shared_bus->children().size();

            if (!shared_bus->remove_child(*device))
            {
              ++failure_count;
            };
          }
        });
      }

      threads.clear();

      THEN("every attach was matched by a successful removal")
      {
        REQUIRE(failure_count == 0);
      }

      THEN("the bus has no children")
      {
        REQUIRE(shared_bus->children().empty());
      }
    }
  }
}

SCENARIO("Concurrent removal of a bus with its own children is race-free")
{
  GIVEN("A shared parent bus, and a pool of child busses attached to it")
  {
    auto parent_bus = kstd::make_shared<kapi::devices::bus>("stress_test_nested_parent");
    kapi::devices::get_root_bus()->add_child(parent_bus);

    constexpr auto pool_size = 8uz;

    auto pool_lock = std::mutex{};
    auto pool = std::vector<kstd::shared_ptr<kapi::devices::bus>>{};

    for (auto i = 0uz; i < pool_size; ++i)
    {
      auto child = kstd::make_shared<kapi::devices::bus>(kstd::format("stress_test_nested_child_{}", i));
      parent_bus->add_child(child);
      pool.push_back(child);
    }

    WHEN("threads concurrently attach devices to shared child busses, while others tear an entire child bus down "
         "through the parent and replace it")
    {
      auto threads = std::vector<std::jthread>{};
      threads.reserve(thread_count);

      for (auto thread_index = 0uz; thread_index < thread_count; ++thread_index)
      {
        threads.emplace_back([&pool, &pool_lock, &parent_bus, thread_index] {
          for (auto i = 0uz; i < devices_per_thread; ++i)
          {
            auto pool_index = (thread_index + i) % pool_size;

            auto child_bus = kstd::shared_ptr<kapi::devices::bus>{};
            {
              auto guard = std::lock_guard{pool_lock};
              child_bus = pool[pool_index];
            }

            auto name = kstd::format("stress_test_nested_device_{}_{}", thread_index, i);
            auto device = kstd::make_shared<test_device>(name);

            child_bus->add_child(device);

            std::ignore = child_bus->remove_child(*device);

            if (i % 25 == 0)
            {
              if (parent_bus->remove_child(*child_bus))
              {
                auto replacement =
                    kstd::make_shared<kapi::devices::bus>(kstd::format("stress_test_nested_child_{}", pool_index));
                parent_bus->add_child(replacement);

                auto guard = std::lock_guard{pool_lock};
                pool[pool_index] = replacement;
              }
            }
          }
        });
      }

      threads.clear();

      THEN("nothing crashed or deadlocked")
      {
        SUCCEED();
      }
    }
  }
}