#include <kapi/memory.hpp>

#include <kernel/memory/bitmap_allocator.hpp>
#include <kernel/memory/mmio_allocator.hpp>

#include <kapi/system.hpp>

#include <kstd/print.hpp>
#include <kstd/units.hpp>

#include <algorithm>
#include <cstddef>
#include <cstdint>
#include <optional>
#include <ranges>
#include <span>
#include <utility>

namespace kapi::memory
{

  namespace
  {
    struct bad_frame_allocator final : public frame_allocator
    {
      bad_frame_allocator static instance;

      auto allocate_many(std::size_t) noexcept -> std::optional<std::pair<frame, std::size_t>> override
      {
        system::panic("Tried to allocate frames without an active allocator.");
      }

      auto mark_used(frame) -> void override
      {
        system::panic("Tried to mark frame as used without an active allocator.");
      }

      auto release_many(std::pair<frame, std::size_t>) -> void override
      {
        system::panic("Tried to release frames without an active allocator.");
      }
    };

    struct bad_page_mapper final : public page_mapper
    {
      bad_page_mapper static instance;

      auto map(page, frame, flags) -> std::byte * override
      {
        system::panic("Tried to map a page without an active mapper.");
      }

      auto unmap(page) -> void override
      {
        system::panic("Tried to unmap a page without an active mapper.");
      }

      auto try_unmap(page) noexcept -> bool override
      {
        return false;
      }
    };

    constinit bad_frame_allocator bad_frame_allocator::instance{};
    constinit bad_page_mapper bad_page_mapper::instance{};
    auto constinit allocator = std::optional<kernel::memory::bitmap_frame_allocator>{};
    auto constinit mmio_allocator = std::optional<kernel::memory::mmio_allocator>{};
  }  // namespace

  constinit auto static active_frame_allocator = static_cast<frame_allocator *>(&bad_frame_allocator::instance);
  constinit auto static active_page_mapper = static_cast<page_mapper *>(&bad_page_mapper::instance);

  auto get_frame_allocator() -> frame_allocator &
  {
    return *active_frame_allocator;
  }

  auto set_frame_allocator(frame_allocator & allocator) -> std::optional<frame_allocator *>
  {
    if (&allocator == active_frame_allocator)
    {
      return {};
    }
    return std::exchange(active_frame_allocator, &allocator);
  }

  auto set_page_mapper(page_mapper & mapper) -> std::optional<page_mapper *>
  {
    if (&mapper == active_page_mapper)
    {
      return {};
    }
    return std::exchange(active_page_mapper, &mapper);
  }

  auto allocate_frame() -> std::optional<frame>
  {
    return get_frame_allocator().allocate();
  }

  auto allocate_many_frames(std::size_t count) -> std::optional<std::pair<frame, std::size_t>>
  {
    return get_frame_allocator().allocate_many(count);
  }

  auto map(page page, frame frame, page_mapper::flags flags) -> std::byte *
  {
    return active_page_mapper->map(page, frame, flags);
  }

  auto unmap(page page) -> void
  {
    return active_page_mapper->unmap(page);
  }

  auto init_pmm(std::size_t frame_count, void (&handoff_handler)(frame_allocator &)) -> void
  {
    using namespace kstd::units_literals;

    auto const bitmap_bytes = kstd::units::bytes{(frame_count + 7uz) / 8uz};
    auto const bitmap_pages = (bitmap_bytes + page::size - 1_B) / page::size;

    auto const bitmap_frames = allocate_many_frames(bitmap_pages);
    if (!bitmap_frames)
    {
      system::panic("[OS:MEM] Not enough memory for bitmap allocator!");
    }

    auto const flags = page_mapper::flags::writable | page_mapper::flags::supervisor_only | page_mapper::flags::global;
    auto bitmap_ptr = static_cast<std::uint64_t *>(nullptr);

    std::ranges::for_each(std::views::iota(0uz, bitmap_pages), [&](auto index) {
      auto page = page::containing(pmm_metadata_base + index * page::size);
      auto frame = memory::frame(bitmap_frames->first.number() + index);
      auto mapped = active_page_mapper->map(page, frame, flags);
      if (!bitmap_ptr)
      {
        bitmap_ptr = reinterpret_cast<std::uint64_t *>(mapped);
      }
    });

    auto bitmap =
        std::span{bitmap_ptr, (bitmap_bytes + kstd::type_size<std::uint64_t> - 1_B) / kstd::type_size<std::uint64_t>};

    allocator.emplace(bitmap, frame_count);

    handoff_handler(allocator.value());
    set_frame_allocator(allocator.value());
    kstd::println("[OS:MEM] Physical memory manager initialized.");
  }

  auto init_mmio(linear_address base, std::size_t page_count) -> void
  {
    mmio_allocator.emplace(base, page_count);
  }

  auto allocate_mmio_region(std::size_t page_count) -> mmio_region
  {
    auto region = mmio_allocator->allocate(page_count);
    return {region, page_count};
  }

  auto map_mmio_region(mmio_region region, physical_address hw_base, page_mapper::flags flags) -> std::byte *
  {
    auto start_page = page::containing(region.first);
    auto start_frame = frame::containing(hw_base);

    flags |= page_mapper::flags::uncached;

    auto start = map(start_page, start_frame, flags);

    std::ranges::for_each(std::views::iota(1uz, region.second), [&](auto index) {
      auto page = page::containing(region.first + index * page::size);
      auto frame = frame::containing(hw_base + index * page::size);
      map(page, frame, flags);
    });

    return start;
  }

  auto release_mmio_region(mmio_region region) -> void
  {
    mmio_allocator->release(region.first);
  }

}  // namespace kapi::memory