#include "kernel/memory/bitmap_allocator.hpp"

#include "kapi/memory.hpp"

#include "catch2/matchers/catch_matchers.hpp"
#include "catch2/matchers/catch_matchers_range_equals.hpp"

#include <catch2/catch_test_macros.hpp>

#include <cstdint>
#include <limits>
#include <ranges>
#include <span>
#include <vector>

constexpr auto all_bits_set = std::numeric_limits<std::uint64_t>::max();
constexpr auto available_frames = 1024uz;

SCENARIO("Bitmap allocator construction and initialization", "[memory][bitmap_allocator]")
{
  GIVEN("A storage region")
  {
    auto storage = std::vector(available_frames / 64, 0uz);

    WHEN("constructing the allocator with 0 frames")
    {
      auto allocator = kernel::memory::bitmap_frame_allocator{std::span(storage), 0};

      THEN("the storage region is not modified")
      {
        REQUIRE_THAT(storage, Catch::Matchers::RangeEquals(std::vector(16, 0uz)));
      }
    }

    WHEN("constructing with 1 frame")
    {
      auto allocator = kernel::memory::bitmap_frame_allocator{std::span(storage), 1};

      THEN("the first word of the storage region is set to all ones")
      {
        REQUIRE_THAT(std::views::take(storage, 1), Catch::Matchers::RangeEquals(std::vector(1, all_bits_set)));
      }

      THEN("the rest of the storage region is not modified")
      {
        REQUIRE_THAT(std::views::drop(storage, 1), Catch::Matchers::RangeEquals(std::vector(15, 0uz)));
      }
    }

    WHEN("constructing with 64 frames")
    {
      auto allocator = kernel::memory::bitmap_frame_allocator{std::span(storage), 64};

      THEN("the first word of the storage region is set to all ones")
      {
        REQUIRE_THAT(std::views::take(storage, 1), Catch::Matchers::RangeEquals(std::vector(1, all_bits_set)));
      }

      THEN("the rest of the storage region is not modified")
      {
        REQUIRE_THAT(std::views::drop(storage, 1), Catch::Matchers::RangeEquals(std::vector(15, 0uz)));
      }
    }

    WHEN("constructing with all available frames")
    {
      auto allocator = kernel::memory::bitmap_frame_allocator{std::span(storage), available_frames};

      THEN("the storage region is filled with all ones")
      {
        REQUIRE_THAT(storage, Catch::Matchers::RangeEquals(std::vector(16, all_bits_set)));
      }
    }

    WHEN("constructing with half the available frames")
    {
      auto allocator = kernel::memory::bitmap_frame_allocator{std::span(storage), available_frames / 2};

      THEN("the first half of the storage region is filled with all ones")
      {
        REQUIRE_THAT(std::views::take(storage, (available_frames / 2) / 64),
                     Catch::Matchers::RangeEquals(std::vector((available_frames / 2) / 64, all_bits_set)));
      }

      THEN("the second half of the storage region is filled with all zeros")
      {
        REQUIRE_THAT(std::views::drop(storage, (available_frames / 2) / 64),
                     Catch::Matchers::RangeEquals(std::vector((available_frames / 2) / 64, 0uz)));
      }
    }
  }
}

SCENARIO("Bitmap allocator frame allocation", "[memory][bitmap_allocator]")
{
  GIVEN("A storage region")
  {
    auto storage = std::vector(available_frames / 64, 0uz);

    AND_GIVEN("an allocator constructed with all available frames but no free ones")
    {
      auto allocator = kernel::memory::bitmap_frame_allocator{std::span(storage), available_frames};

      WHEN("allocating 1 frame")
      {
        auto result = allocator.allocate_many(1);

        THEN("the result is empty")
        {
          REQUIRE_FALSE(result.has_value());
        }
      }
    }

    AND_GIVEN("an allocator constructed with all available frames but only one free one")
    {
      auto allocator = kernel::memory::bitmap_frame_allocator{std::span(storage), available_frames};
      allocator.release_many({kapi::memory::frame{0}, 1});

      WHEN("allocating 1 frame")
      {
        auto result = allocator.allocate_many(1);

        THEN("the result is not empty")
        {
          REQUIRE(result.has_value());
        }

        THEN("the result contains 1 frame")
        {
          REQUIRE(result->second == 1);
        }
      }

      WHEN("allocating more frames than are free")
      {
        auto result = allocator.allocate_many(2);

        THEN("the result is empty")
        {
          REQUIRE_FALSE(result.has_value());
        }

        AND_WHEN("allocating a single frame")
        {
          auto result = allocator.allocate_many(1);

          THEN("the result is not empty")
          {
            REQUIRE(result.has_value());
          }
        }

        WHEN("allocating 0 frames")
        {
          auto result = allocator.allocate_many(0);

          THEN("the result is empty")
          {
            REQUIRE_FALSE(result.has_value());
          }
        }
      }
    }

    AND_GIVEN("an allocator with many single frame holes")
    {
      auto allocator = kernel::memory::bitmap_frame_allocator{std::span(storage), available_frames};
      for (auto i = 0uz; i < available_frames; i += 2)
      {
        allocator.release_many({kapi::memory::frame{i}, 1});
      }

      WHEN("allocating 1 frame")
      {
        auto result = allocator.allocate_many(1);

        THEN("the result is not empty")
        {
          REQUIRE(result.has_value());
        }

        THEN("the result contains 1 frame")
        {
          REQUIRE(result->second == 1);
        }
      }

      WHEN("allocating 2 frames")
      {
        auto result = allocator.allocate_many(2);

        THEN("the result is empty")
        {
          REQUIRE_FALSE(result.has_value());
        }
      }
    }

    AND_GIVEN("and allocator with all frames marked as free")
    {
      auto allocator = kernel::memory::bitmap_frame_allocator{std::span(storage), available_frames};
      allocator.release_many({kapi::memory::frame{0}, available_frames});

      WHEN("allocating 1 frame")
      {
        auto result = allocator.allocate_many(1);

        THEN("the result is not empty")
        {
          REQUIRE(result.has_value());
        }

        THEN("the result contains 1 frame")
        {
          REQUIRE(result->second == 1);
        }
      }

      WHEN("allocating multiple frames")
      {
        auto result = allocator.allocate_many(20);

        THEN("the result is not empty")
        {
          REQUIRE(result.has_value());
        }

        THEN("the result contains 20 frames")
        {
          REQUIRE(result->second == 20);
        }
      }

      WHEN("marking all frames as used")
      {
        for (auto i = 0uz; i < available_frames; i++)
        {
          allocator.mark_used(kapi::memory::frame{i});
        }

        THEN("the allocator has no free frames")
        {
          REQUIRE_FALSE(allocator.allocate_many(1).has_value());
        }
      }
    }

    AND_GIVEN("an allocator with a contiguous block of free frames")
    {
      auto allocator = kernel::memory::bitmap_frame_allocator{std::span(storage), available_frames};
      allocator.release_many({kapi::memory::frame{0}, available_frames});
      for (auto i = 0uz; i < available_frames / 2; i++)
      {
        allocator.mark_used(kapi::memory::frame{i});
      }

      WHEN("allocating a single frame")
      {
        auto result = allocator.allocate_many(1);

        THEN("the result is not empty")
        {
          REQUIRE(result.has_value());
        }

        THEN("the result contains 1 frame")
        {
          REQUIRE(result->second == 1);
        }
      }

      WHEN("allocating multiple frames")
      {
        auto result = allocator.allocate_many(20);

        THEN("the result is not empty")
        {
          REQUIRE(result.has_value());
        }

        THEN("the result contains 20 frames")
        {
          REQUIRE(result->second == 20);
        }
      }
    }
  }
}