#include <kapi/devices/driver_registry.hpp>

#include <kapi/devices.hpp>

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

#include <catch2/catch_test_macros.hpp>

#include <cstdint>
#include <string_view>

namespace
{

  struct test_claim
  {
    constexpr auto static id = kapi::capabilities::facet_id{"test_identification"};

    virtual ~test_claim() = 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
  {
    constexpr auto static id = kapi::capabilities::facet_id{"test_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.facet<test_claim>();
      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
  {
    test_bus(kstd::string name)
        : bus{name, 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_claim
  {
    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;
    }

    auto name() const noexcept -> std::string_view override
    {
      return "Driver Registry Test Device";
    }

    mutable unsigned probe_calls;

  protected:
    auto query_facet(kapi::capabilities::facet_id facet) -> void * override
    {
      if (facet == test_claim::id)
      {
        return static_cast<test_claim *>(this);
      }

      return kapi::devices::driver::query_facet(facet);
    }

  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());
      }
    }
  }
}

SCENARIO("Driver registry leaves device in failed state if driver probe fails", "[devices][driver_registry]")
{
  GIVEN("A driver that fails to probe")
  {
    auto bus = kstd::make_shared<test_bus>("driver_registry_tie_bus");
    auto driver = kstd::make_shared<test_driver>(3u, false);

    kapi::devices::driver_registry::get().add(driver);

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

      THEN("the device ends up in failed state")
      {
        auto dev = kapi::devices::device_registry::get().find("driver_probe_failure_device");
        REQUIRE(dev != nullptr);
        REQUIRE(dev->state() == kapi::devices::state::failed);
      }
    }
  }
}

SCENARIO("Driver registry leaves bound driver bound, even if better driver arrives", "[devices][driver_registry]")
{
  GIVEN("a low priority driver")
  {
    auto bus = kstd::make_shared<test_bus>("driver_registry_tie_bus");
    auto lower = kstd::make_shared<test_driver>(0, true);

    kapi::devices::driver_registry::get().add(lower);

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

      AND_WHEN("a second, higher priority driver arrives")
      {
        auto higher = kstd::make_shared<test_driver>(0, true);
        kapi::devices::driver_registry::get().add(higher);

        THEN("the device is still bound to the lower priority driver")
        {
          auto dev = kapi::devices::device_registry::get().find("driver_probe_failure_device");
          REQUIRE(dev != nullptr);
          REQUIRE(dev->state() == kapi::devices::state::bound);
          REQUIRE(dev->bound_driver() == lower.get());
        }
      }
    }
  }
}