#include "kernel/memory/mmio_allocator.hpp"

#include "kapi/memory.hpp"
#include "kapi/system.hpp"

#include <kstd/allocator>

#include <cstddef>
#include <memory>
#include <utility>

namespace kernel::memory
{

  mmio_allocator::mmio_allocator(kapi::memory::linear_address base, std::size_t pages)
      : m_head{make_node(base, pages, nullptr, nullptr, true)}
  {}

  auto mmio_allocator::allocate(std::size_t count) -> kapi::memory::linear_address
  {
    if (count == 0 || !m_head)
    {
      return {};
    }

    auto current = m_head;
    while (current)
    {
      if (current->is_free && current->page_count >= count)
      {
        if (current->page_count > count)
        {
          auto new_base = current->base + (count * kapi::memory::page::size);
          auto split_node = make_node(new_base, current->page_count - count, std::move(current->next), current, true);

          if (current->next)
          {
            current->next->previous = split_node;
          }
          current->next = split_node;
          current->page_count = count;
        }

        current->is_free = false;
        return current->base;
      }
      current = current->next;
    }

    kapi::system::panic("[OS:MEM] MMIO alloctor out of memory!");
    return {};
  }

  auto mmio_allocator::release(kapi::memory::linear_address base) -> void
  {
    auto current = m_head;

    while (current)
    {
      if (current->base == base && !current->is_free)
      {
        current->is_free = true;

        if (current->next && current->next->is_free)
        {
          auto removed = current->next;
          current->page_count += removed->page_count;
          current->next = removed->next;
          if (current->next)
          {
            current->next->previous = current;
          }
          destroy_node(removed);
        }

        if (current->previous && current->previous->is_free)
        {
          auto removed = current;
          removed->previous->page_count += removed->page_count;
          removed->previous->next = removed->next;
          if (removed->next)
          {
            removed->next->previous = removed->previous;
          }
          destroy_node(removed);
        }
        return;
      }
      current = current->next;
    }
  }

  auto mmio_allocator::make_node(kapi::memory::linear_address base, std::size_t page_count, node * next,
                                 node * previous, bool is_free) -> node *
  {
    using traits = std::allocator_traits<kstd::allocator<node>>;

    auto new_node = traits::allocate(m_allocator, 1);
    traits::construct(m_allocator, new_node, base, page_count, next, previous, is_free);
    return new_node;
  }

  auto mmio_allocator::destroy_node(node * instance) -> void
  {
    using traits = std::allocator_traits<kstd::allocator<node>>;

    traits::destroy(m_allocator, instance);
    traits::deallocate(m_allocator, instance, 1);
  }

}  // namespace kernel::memory