#include <kapi/devices/driver_registry.hpp>

#include <kapi/devices.hpp>

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

#include <catch2/catch_test_macros.hpp>

#include <cstdint>

namespace
{

  struct test_identification
  {
    constexpr auto static id = kapi::devices::interface{"test_identification"};

    virtual ~test_identification() = default;

    [[nodiscard]] virtual auto priority() const noexcept -> std::uint32_t = 0;
    [[nodiscard]] virtual auto probe_succeeds() const noexcept -> bool = 0;
  };

  struct test_protocol final : kapi::devices::bus_protocol
  {
    explicit test_protocol(unsigned * match_calls)
        : match_calls{match_calls}
    {}

    auto enumerate(kapi::devices::bus &) -> void override {}

    [[nodiscard]] auto match(kapi::devices::device const &, kapi::devices::driver const & driver) const
        -> kstd::result<std::uint32_t> override
    {
      ++*match_calls;

      auto const * identification = driver.as<test_identification>();
      if (!identification)
      {
        return kstd::failure(kapi::devices::driver_match_errc::no_match);
      }

      return identification->priority();
    }

    unsigned * match_calls{};
  };

  struct test_bus final : kapi::devices::bus
  {
    using kapi::devices::bus::bus;

    [[nodiscard]] auto protocol() -> kapi::devices::bus_protocol * override
    {
      return &bus_protocol;
    }

    unsigned match_calls{};
    test_protocol bus_protocol{&match_calls};
  };

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

  struct test_driver final : kapi::devices::driver, test_identification
  {
    test_driver(std::uint32_t priority, bool probe_succeeds)
        : probe_calls{}
        , m_priority{priority}
        , m_probe_succeeds{probe_succeeds}
    {}

    [[nodiscard]] auto probe(kapi::devices::device &) -> kstd::result<void> override
    {
      ++probe_calls;

      if (!m_probe_succeeds)
      {
        return kstd::failure(kapi::devices::driver_match_errc::no_match);
      }

      return kstd::success();
    }

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

    auto priority() const noexcept -> std::uint32_t override
    {
      return m_priority;
    }

    auto probe_succeeds() const noexcept -> bool override
    {
      return m_probe_succeeds;
    }

    mutable unsigned probe_calls;

  protected:
    auto query_interface(kapi::devices::interface interface) -> void * override
    {
      if (interface == test_identification::id)
      {
        return static_cast<test_identification *>(this);
      }

      return kapi::devices::driver::query_interface(interface);
    }

  private:
    std::uint32_t m_priority;
    bool m_probe_succeeds;
  };

}  // namespace

SCENARIO("Driver registry binds devices and drivers regardless of arrival order", "[devices][driver_registry]")
{
  GIVEN("a bus and a driver that always matches and probes successfully")
  {
    auto bus = kstd::make_shared<test_bus>("driver_registry_order_bus");
    auto driver = kstd::make_shared<test_driver>(1u, true);

    WHEN("the driver registers before the device arrives")
    {
      kapi::devices::driver_registry::get().add(driver);
      bus->add_child(kstd::make_shared<test_device>("driver_registry_test_device_1"));

      THEN("the device binds")
      {
        auto dev = kapi::devices::device_registry::get().find("driver_registry_test_device_1");
        REQUIRE(dev);
        REQUIRE(dev->state() == kapi::devices::state::bound);
      }
    }

    WHEN("the device arrives before the driver registers")
    {
      bus->add_child(kstd::make_shared<test_device>("driver_registry_test_device_2"));
      kapi::devices::driver_registry::get().add(driver);

      THEN("the device binds")
      {
        auto dev = kapi::devices::device_registry::get().find("driver_registry_test_device_2");
        REQUIRE(dev);
        REQUIRE(dev->state() == kapi::devices::state::bound);
      }
    }
  }
}

SCENARIO("driver_registry picks the highest-priority match", "[devices][driver_registry]")
{
  GIVEN("two drivers with different priorities")
  {
    auto bus = kstd::make_shared<test_bus>("driver_registry_priority_bus");
    auto low_priority = kstd::make_shared<test_driver>(1u, true);
    auto high_priority = kstd::make_shared<test_driver>(5u, true);

    kapi::devices::driver_registry::get().add(low_priority);
    kapi::devices::driver_registry::get().add(high_priority);

    WHEN("the device arrives")
    {
      bus->add_child(kstd::make_shared<test_device>("driver_registry_priority_device"));

      THEN("it binds to the higher-priority driver")
      {
        auto dev = kapi::devices::device_registry::get().find("driver_registry_priority_device");
        REQUIRE(dev != nullptr);
        REQUIRE(dev->bound_driver() == high_priority.get());
      }
    }
  }
}

SCENARIO("driver_registry breaks ties by registration order", "[devices][driver_registry]")
{
  GIVEN("two drivers with equal priority, registered in a specific order")
  {
    auto bus = kstd::make_shared<test_bus>("driver_registry_tie_bus");
    auto registered_first = kstd::make_shared<test_driver>(3u, true);
    auto registered_second = kstd::make_shared<test_driver>(3u, true);

    kapi::devices::driver_registry::get().add(registered_first);
    kapi::devices::driver_registry::get().add(registered_second);

    WHEN("the device arrives")
    {
      bus->add_child(kstd::make_shared<test_device>("driver_registry_tie_device"));

      THEN("the first-registered driver wins the tie")
      {
        auto dev = kapi::devices::device_registry::get().find("driver_registry_tie_device");
        REQUIRE(dev != nullptr);
        REQUIRE(dev->bound_driver() == registered_first.get());
      }
    }
  }
}