#include <kapi/devices.hpp>

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

#include <catch2/catch_test_macros.hpp>

#include <cstdint>
#include <string_view>
#include <thread>
#include <tuple>
#include <vector>

namespace
{
  struct test_protocol final : kapi::devices::bus_protocol
  {
    auto enumerate(kapi::devices::bus &) -> void override {}

    [[nodiscard]] auto match(kapi::devices::device const &, kapi::devices::driver const &) const
        -> kstd::result<std::uint32_t> override
    {
      return 0u;
    }
  } constinit test_protocol_instance{};

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

  struct test_driver final : kapi::devices::driver
  {
    [[nodiscard]] auto probe(kapi::devices::device &) -> kstd::result<void> override
    {
      return kstd::success();
    }

    auto unbind(kapi::devices::device &) -> void override {}

    [[nodiscard]] auto name() const noexcept -> std::string_view override
    {
      return "driver_registry_stress_test_driver";
    }
  };

}  // namespace

constexpr auto driver_thread_count = 4;
constexpr auto device_thread_count = 4;
constexpr auto drivers_per_thread = 100;
constexpr auto devices_per_thread = 200;

SCENARIO("Concurrent driver registration and device attachment is race-free")
{
  GIVEN("A bus using a protocol that matches every device, shared by several threads")
  {
    auto shared_bus = kstd::make_shared<kapi::devices::bus>("driver_registry_stress_test_bus", test_protocol_instance);
    kapi::devices::get_root_bus()->add_child(shared_bus);

    WHEN("some threads concurrently register new drivers while others attach and detach devices")
    {
      auto threads = std::vector<std::jthread>{};
      threads.reserve(driver_thread_count + device_thread_count);

      for (auto i = 0; i < driver_thread_count; ++i)
      {
        threads.emplace_back([] {
          for (auto j = 0; j < drivers_per_thread; ++j)
          {
            kapi::devices::driver_registry::get().add(kstd::make_shared<test_driver>());
          }
        });
      }

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

            shared_bus->add_child(device);
            std::ignore = shared_bus->remove_child(*device);
          }
        });
      }

      threads.clear();

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