#include <kernel/memory/block_list_allocator.hpp>

#include <kernel/memory/heap_allocator.hpp>

#include <kapi/memory.hpp>
#include <kapi/system.hpp>

#include <kstd/mutex>
#include <kstd/units>

#include <bit>
#include <cstddef>
#include <cstdint>
#include <memory>

using namespace kstd::units_literals;

namespace kernel::memory
{

  namespace
  {
    [[nodiscard]] constexpr auto align_up(std::byte * pointer, kstd::units::bytes alignment) noexcept -> std::byte *
    {
      auto const remainder = std::bit_cast<std::uintptr_t>(pointer) % static_cast<std::size_t>(alignment);
      return remainder == 0 ? pointer : pointer + static_cast<std::size_t>(alignment) - remainder;
    }
  }  // namespace

  block_list_allocator::block_list_allocator(kapi::memory::linear_address base) noexcept
      : heap_allocator{}
      , m_base{base}
      , m_frontier{base}
      , m_block_list{}
      , m_lock{}
  {}

  auto block_list_allocator::allocate(kstd::units::bytes size, kstd::units::bytes alignment) noexcept -> void *
  {
    kstd::lock_guard guard{m_lock};

    for (auto attempt = 0uz; attempt < 2uz; ++attempt)
    {
      auto current = m_block_list;
      while (current != nullptr)
      {
        if (current->free)
        {
          auto const raw_block = reinterpret_cast<std::byte *>(current);
          auto const unaligned_payload = raw_block + allocated_metadata_size;
          auto const aligned_payload = align_up(unaligned_payload, alignment);
          auto const required_padding = static_cast<kstd::units::bytes>(aligned_payload - unaligned_payload);
          auto const total_required_size = required_padding + allocated_metadata_size + size;

          if (current->usable_size >= total_required_size)
          {
            auto const payload_header = aligned_payload - sizeof(block_header *) - sizeof(block_header);
            auto const front_padding = static_cast<kstd::units::bytes>(payload_header - raw_block);

            auto payload_block = current;

            if (front_padding >= allocated_metadata_size + minimum_allocation_size)
            {
              payload_block = reinterpret_cast<block_header *>(payload_header);
              std::construct_at(payload_block, current->usable_size - front_padding, true, current->next);

              if (payload_block->next)
              {
                payload_block->next->prev = payload_block;
              }

              current->usable_size = front_padding - allocated_metadata_size;
              current->next = payload_block;
              payload_block->prev = current;
            }

            auto const header_size =
                static_cast<kstd::units::bytes>(aligned_payload - reinterpret_cast<std::byte *>(payload_block));
            auto const payload_size = header_size - allocated_metadata_size + size;
            split(payload_block, payload_size, 0_B);

            payload_block->free = false;

            auto back_pointer = reinterpret_cast<block_header **>(aligned_payload - sizeof(block_header *));
            *back_pointer = payload_block;

            return reinterpret_cast<void *>(aligned_payload);
          }
        }
        current = current->next;
      }

      auto const search_size = size + alignment;
      if (attempt == 0uz && !expand(search_size))
      {
        return nullptr;
      }
    }

    return nullptr;
  }

  auto block_list_allocator::deallocate(void * ptr) noexcept -> void
  {
    if (!ptr)
    {
      return;
    }

    kstd::lock_guard guard{m_lock};

    auto const aligned_payload = reinterpret_cast<std::byte *>(ptr);
    auto back_pointer = reinterpret_cast<block_header **>(aligned_payload - sizeof(block_header *));
    auto block = *back_pointer;

    block->free = true;
    coalesce(block);
  }

  auto block_list_allocator::expand(kstd::units::bytes size) noexcept -> bool
  {
    auto const total_required_size = size + allocated_metadata_size;
    auto const frames_needed = (total_required_size + kapi::memory::frame::size - 1_B) / kapi::memory::frame::size;

    auto const flags = kapi::memory::page_mapper::flags::writable | kapi::memory::page_mapper::flags::supervisor_only |
                       kapi::memory::page_mapper::flags::global;

    for (auto i = 0uz; i < frames_needed; ++i)
    {
      auto frame = kapi::memory::allocate_frame();
      if (!frame)
      {
        kapi::system::panic("[OS:Heap] OOM when expanding heap.");
        return false;
      }

      auto page = kapi::memory::page::containing(m_frontier + i * kapi::memory::page::size);
      kapi::memory::map(page, *frame, flags);
    }

    auto block = static_cast<block_header *>(m_frontier);
    std::construct_at(block, frames_needed * kapi::memory::frame::size - allocated_metadata_size, true, nullptr,
                      nullptr);

    m_frontier += frames_needed * kapi::memory::frame::size;

    if (!m_block_list)
    {
      m_block_list = block;
    }
    else
    {
      auto current = m_block_list;
      while (current->next)
      {
        current = current->next;
      }
      current->next = block;
      block->prev = current;
    }

    coalesce(block);
    return true;
  }

  auto block_list_allocator::coalesce(block_header * block) noexcept -> void
  {
    if (block->next && block->next->free)
    {
      auto follower = block->next;
      block->usable_size += follower->usable_size + allocated_metadata_size;
      block->next = follower->next;
      if (block->next)
      {
        block->next->prev = block;
      }
      std::destroy_at(follower);
    }

    if (block->prev && block->prev->free)
    {
      auto leader = block->prev;
      leader->usable_size += block->usable_size + allocated_metadata_size;
      leader->next = block->next;
      if (block->next)
      {
        block->next->prev = leader;
      }
      std::destroy_at(block);
    }
  }

  auto block_list_allocator::split(block_header * block, kstd::units::bytes size, kstd::units::bytes padding) noexcept
      -> void
  {
    auto const new_block_size = size + padding;

    if (block->usable_size > new_block_size + allocated_metadata_size + minimum_allocation_size)
    {
      auto raw_block = reinterpret_cast<std::byte *>(block);
      auto new_block = reinterpret_cast<block_header *>(raw_block + allocated_metadata_size + new_block_size);
      std::construct_at(new_block, block->usable_size - new_block_size - allocated_metadata_size, true, block->next,
                        block);

      if (new_block->next)
      {
        new_block->next->prev = new_block;
      }

      block->usable_size = new_block_size;
      block->next = new_block;
    }
  }

}  // namespace kernel::memory