#include <kapi/devices/driver_registry.hpp>

#include <kapi/devices/bus.hpp>
#include <kapi/devices/bus_protocol.hpp>
#include <kapi/devices/device.hpp>
#include <kapi/devices/driver.hpp>
#include <kapi/devices/manager.hpp>
#include <kapi/system.hpp>
#include <kapi/test_support/devices.hpp>

#include <kstd/memory.hpp>
#include <kstd/print.hpp>
#include <kstd/system_error.hpp>
#include <kstd/vector.hpp>

#include <algorithm>
#include <cstdint>
#include <functional>
#include <optional>
#include <utility>

namespace kapi::devices
{

  namespace
  {
    constinit auto static registry = std::optional<driver_registry>{};

    struct candidate
    {
      std::uint32_t priority;
      kstd::shared_ptr<driver> driver_handle;
    };
  }  // namespace

  auto driver_registry::init() -> void
  {
    if (registry.has_value())
    {
      system::panic("[kernel] Device driver registry has already been initialized.");
    }

    registry.emplace();
  }

  auto driver_registry::get() -> driver_registry &
  {
    if (!registry)
    {
      system::panic("[kernel] Device driver registry has not been initialized.");
    }

    return *registry;
  }

  auto driver_registry::add(kstd::shared_ptr<driver> driver) -> void
  {
    m_drivers.emplace_back(std::move(driver));

    for (auto const & device : kapi::devices::all())
    {
      if (device->state() != state::bound)
      {
        try_bind(device);
      }
    }
  }

  auto driver_registry::device_attached(kstd::shared_ptr<device> const & device) -> void
  {
    try_bind(device);
  }

  auto driver_registry::try_bind(kstd::shared_ptr<device> const & device) -> void
  {
    if (!device || device->state() == state::bound)
    {
      return;
    }

    auto parent = device->parent();
    auto protocol = parent ? parent->protocol() : nullptr;
    if (!protocol)
    {
      return;
    }

    auto candidates = kstd::vector<candidate>{};

    for (auto const & driver : m_drivers)
    {
      auto match = protocol->match(*device, *driver);
      if (match)
      {
        candidates.push_back({.priority = *match, .driver_handle = driver});
      }
      else if (match.error() != kstd::errc::not_supported)
      {
        kstd::println(kstd::print_sink::stderr, "[OS:DRV] match() failed for device {}: {}", device->name(),
                      match.error().message());
      }
    }

    std::ranges::stable_sort(candidates, std::ranges::greater{}, &candidate::priority);

    for (auto const & candidate : candidates)
    {
      auto probed = candidate.driver_handle->probe(*device);
      if (probed)
      {
        device->bind_driver(candidate.driver_handle);
        device->set_state(state::bound);
        kstd::println("[OS:DRV] Bound device {} (priority {})", device->name(), candidate.priority);
        return;
      }

      kstd::println(kstd::print_sink::stderr, "[OS:DRV] probe() failed for device {} (priority {}): {}", device->name(),
                    candidate.priority, probed.error().message());
    }

    if (!candidates.empty())
    {
      device->set_state(state::failed);
    }
  }

}  // namespace kapi::devices

namespace kapi::test_support::devices
{
  auto deinit_driver_registry() -> void
  {
    kapi::devices::registry.reset();
  }
}  // namespace kapi::test_support::devices