#include <kapi/devices/power.hpp>

#include <kapi/devices.hpp>

#include <kstd/format.hpp>
#include <kstd/memory.hpp>
#include <kstd/result.hpp>
#include <kstd/string.hpp>
#include <kstd/system_error.hpp>
#include <kstd/vector.hpp>

#include <catch2/catch_test_macros.hpp>

#include <algorithm>
#include <cstddef>
#include <string_view>

namespace
{

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

  struct recording_driver final : kapi::devices::driver
  {
    explicit recording_driver(kstd::vector<kstd::string> & log, kstd::string fail_on_suspend = {})
        : m_log{&log}
        , m_fail_on_suspend{fail_on_suspend}
    {}

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

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

    [[nodiscard]] auto suspend(kapi::devices::device & device) -> kstd::result<void> override
    {
      if (!m_fail_on_suspend.empty() && device.name() == m_fail_on_suspend)
      {
        return kstd::failure(make_error_code(kstd::errc::io_error));
      }

      m_log->push_back(kstd::format("suspend:{}", device.name()));
      return kstd::success();
    }

    [[nodiscard]] auto resume(kapi::devices::device & device) -> kstd::result<void> override
    {
      m_log->push_back(kstd::format("resume:{}", device.name()));
      return kstd::success();
    }

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

  private:
    kstd::vector<kstd::string> * m_log;
    kstd::string m_fail_on_suspend;
  };

  auto index_of(kstd::vector<kstd::string> const & log, std::string_view entry) -> std::size_t
  {
    auto found = std::ranges::find(log, entry);
    REQUIRE(found != log.end());
    return static_cast<std::size_t>(found - log.begin());
  }

}  // namespace

SCENARIO("Suspension and resumption happens in the correct order", "[devices][power]")
{
  GIVEN("a bus with two bound lead devices, itself bound to its own driver, attached to a root")
  {
    auto log = kstd::vector<kstd::string>{};

    auto root = kstd::make_shared<kapi::devices::bus>("power_test_root");
    auto middle = kstd::make_shared<kapi::devices::bus>("power_test_middle");
    auto leaf_1 = kstd::make_shared<test_device>("power_test_leaf_1");
    auto leaf_2 = kstd::make_shared<test_device>("power_test_leaf_2");

    auto middle_driver = kstd::make_shared<recording_driver>(log);
    auto leaf_1_driver = kstd::make_shared<recording_driver>(log);
    auto leaf_2_driver = kstd::make_shared<recording_driver>(log);

    root->add_child(middle);
    middle->add_child(leaf_1);
    middle->add_child(leaf_2);

    middle->bind_driver(middle_driver);
    middle->set_state(kapi::devices::state::bound);

    leaf_1->bind_driver(leaf_1_driver);
    leaf_1->set_state(kapi::devices::state::bound);

    leaf_2->bind_driver(leaf_2_driver);
    leaf_2->set_state(kapi::devices::state::bound);

    WHEN("the bus is suspended")
    {
      auto result = kapi::devices::suspend_tree(*root);

      THEN("it succeeds")
      {
        REQUIRE(result);
      }

      THEN("the leaves were suspended before the middle bus")
      {
        REQUIRE(log.size() == 3);

        auto middle_index = index_of(log, "suspend:power_test_middle");
        auto leaf_1_index = index_of(log, "suspend:power_test_leaf_1");
        auto leaf_2_index = index_of(log, "suspend:power_test_leaf_2");

        REQUIRE(leaf_1_index < middle_index);
        REQUIRE(leaf_2_index < middle_index);
      }

      AND_WHEN("the bus is resumed")
      {
        kapi::devices::resume_tree(*root);

        THEN("the leaves were resumed after the middle bus")
        {
          auto middle_index = index_of(log, "resume:power_test_middle");
          auto leaf_1_index = index_of(log, "resume:power_test_leaf_1");
          auto leaf_2_index = index_of(log, "resume:power_test_leaf_2");

          REQUIRE(middle_index < leaf_1_index);
          REQUIRE(middle_index < leaf_2_index);
        }
      }
    }
  }

  GIVEN("a subtree where one leaf will fail to suspend")
  {
    auto log = kstd::vector<kstd::string>{};

    auto root = kstd::make_shared<kapi::devices::bus>("power_test_root_failing");
    auto middle = kstd::make_shared<kapi::devices::bus>("power_test_middle_failing");
    auto leaf_1 = kstd::make_shared<test_device>("power_test_leaf_1_failing");
    auto leaf_2 = kstd::make_shared<test_device>("power_test_leaf_2_failing");

    auto middle_driver = kstd::make_shared<recording_driver>(log);
    // leaf_1 suspends normally; leaf_2 is told to fail.
    auto leaf_1_driver = kstd::make_shared<recording_driver>(log);
    auto leaf_2_driver = kstd::make_shared<recording_driver>(log, kstd::string{"power_test_leaf_2_failing"});

    root->add_child(middle);
    middle->add_child(leaf_1);
    middle->add_child(leaf_2);

    middle->bind_driver(middle_driver);
    middle->set_state(kapi::devices::state::bound);
    leaf_1->bind_driver(leaf_1_driver);
    leaf_1->set_state(kapi::devices::state::bound);
    leaf_2->bind_driver(leaf_2_driver);
    leaf_2->set_state(kapi::devices::state::bound);

    WHEN("the subtree is suspended")
    {
      auto result = kapi::devices::suspend_tree(*root);

      THEN("the failure is reported")
      {
        REQUIRE_FALSE(result);
      }

      THEN("the leaf that suspended successfully before the failure was resumed again")
      {
        REQUIRE(log.size() == 2);
        REQUIRE(log[0] == "suspend:power_test_leaf_1_failing");
        REQUIRE(log[1] == "resume:power_test_leaf_1_failing");
      }

      THEN("the middle bus, whose child failed to suspend, was never suspended at all")
      {
        REQUIRE(std::ranges::find(log, "suspend:power_test_middle_failing") == log.end());
      }
    }
  }

  GIVEN("a bus with a child that has no bound driver at all")
  {
    auto root = kstd::make_shared<kapi::devices::bus>("power_test_root_unbound");
    auto unbound_child = kstd::make_shared<test_device>("power_test_unbound_child");

    root->add_child(unbound_child);

    THEN("suspending and resuming it does not crash, and simply does nothing for that device")
    {
      REQUIRE(kapi::devices::suspend_tree(*root));
      REQUIRE_NOTHROW(kapi::devices::resume_tree(*root));
    }
  }
}