/*
    * MemoryAllocator.cpp
    * Heap memory allocator
    * Copyright (c) 2025 Daniel Hammer
*/

#include "Heap.hpp"
#include "Memmap.hpp"
#include "HHDM.hpp"

#include <Libraries/Memory.hpp>
#include <Terminal/Terminal.hpp>
#include <Common/Panic.hpp>
#include <Platform/Util.hpp>

#include <Common/Panic.hpp>
#include <CppLib/Vector.hpp>

#include "PageFrameAllocator.hpp"

namespace Memory
{
    HeapAllocator::Header* HeapAllocator::GetHeader(void* block) {
        return (Header*)(block - sizeof(Header));
    }

    size_t HeapAllocator::GetAllocatedBlockSize(void* ptr) {
        Header* header = GetHeader(ptr);
        return header->size;
    }

    // Insert a free block, keeping the list sorted by address and coalescing with
    // any adjacent neighbors. Coalescing keeps the free list compact and prevents
    // fragmentation from accumulating across alloc/free cycles, which is otherwise
    // pathological because Realloc always allocates a fresh block.
    void HeapAllocator::InsertToFreelist(void* ptr, std::size_t size) {
        if (ptr == nullptr || size == 0) return;

        uintptr_t addr = (uintptr_t)ptr;

        Node* prev = &head;
        Node* current = head.next;
        while (current != nullptr && (uintptr_t)current < addr) {
            prev = current;
            current = current->next;
        }

        bool mergedPrev = false;
        if (prev != &head) {
            uintptr_t prevEnd = (uintptr_t)prev + prev->size;
            if (prevEnd == addr) {
                prev->size += size;
                mergedPrev = true;
            }
        }

        if (current != nullptr && addr + size == (uintptr_t)current) {
            if (mergedPrev) {
                prev->size += current->size;
                prev->next = current->next;
            } else {
                Node* newNode = (Node*)addr;
                newNode->size = size + current->size;
                newNode->next = current->next;
                prev->next = newNode;
            }
        } else if (!mergedPrev) {
            Node* newNode = (Node*)addr;
            newNode->size = size;
            newNode->next = current;
            prev->next = newNode;
        }
    }

    void HeapAllocator::InsertPageToFreelist() {
        InsertToFreelist(Memory::g_pfa->Allocate(), 0x1000);
    }

    void HeapAllocator::InsertPagesToFreelist(std::size_t n) {
        auto ptr = Memory::g_pfa->ReallocConsecutive(nullptr, n);
        size_t size = 0x1000 * n;

        InsertToFreelist(ptr, size);
    }

    HeapAllocator::HeapAllocator()
    {
        InsertPagesToFreelist(0x32);
    }

    void* HeapAllocator::Request(size_t size) {
        Lock.Acquire();

        // Round the user request up so allocated blocks are 8-byte aligned. The
        // Header itself is 16 bytes, so the returned pointer is always aligned.
        size_t alignedSize = (size + 7u) & ~(size_t)7u;
        size_t sizeNeeded = alignedSize + sizeof(Header);

    retry:
        Node* current = head.next;
        Node* prev = &head;

        while (current != nullptr) {
            if (current->size >= sizeNeeded) {
                auto locatedBlockSize = current->size;

                // Only split if the leftover region can hold a free-list Node.
                // Otherwise hand out the whole block; storing a sub-Node-sized
                // remainder would write the Node header past the block.
                bool canSplit = (locatedBlockSize >= sizeNeeded + sizeof(Node));

                prev->next = current->next;
                Header* header = (Header*)current;

                header->magic = headerMagic;
                header->size = canSplit ? alignedSize : (locatedBlockSize - sizeof(Header));

                void* block = (void*)((uintptr_t)header + sizeof(Header));

                if (canSplit) {
                    void* rest = (void*)((uintptr_t)header + sizeNeeded);
                    auto newBlockSize = locatedBlockSize - sizeNeeded;
                    InsertToFreelist(rest, newBlockSize);
                }

                Lock.Release();
                return block;
            }

            prev = current;
            current = current->next;
        }

        // No suitable block found -- grow the heap under the lock so
        // InsertToFreelist is protected from concurrent modification.
        // Ensure the new block leaves room for the Header *and* a splittable
        // free-list Node, so even worst-case requests succeed after growth.
        size_t pagesNeeded = (sizeNeeded + sizeof(Node) + 0xFFF) / 0x1000;
        InsertPagesToFreelist(pagesNeeded);
        goto retry;
    }

    void* HeapAllocator::Realloc(void* ptr, size_t size) {
        auto new_block = Request(size);

        if (ptr != nullptr && new_block != nullptr) {
            size_t oldSize = GetAllocatedBlockSize(ptr);
            size_t copySize = (oldSize < size) ? oldSize : size;
            memcpy(new_block, ptr, copySize);
            Free(ptr);
        }

        return new_block;
    }

    void HeapAllocator::Free(void* ptr) {
        if (ptr == nullptr) return;

        Header* header = GetHeader(ptr);

        Lock.Acquire();

        // Validate the magic under the lock so that two concurrent frees of
        // the same pointer can't both pass the check and both insert. The
        // magic is cleared inside the lock for the same reason.
        if (header->magic != headerMagic) {
            Lock.Release();
            Panic("Bad magic in HeapAllocator header", nullptr);
            return;
        }

        size_t size = header->size;
        header->magic = 0;
        InsertToFreelist((void*)header, size + sizeof(Header));

        Lock.Release();
    }
    
    // Traverses the Allocator's linked list for debugging
    void HeapAllocator::Walk() {
        Node* current = {head.next};
        size_t i{0};

        while (current != nullptr) {
            Kt::KernelLogStream(Kt::DEBUG, "HeapAllocator") << base::dec << i << " " << current->size << " bytes & address 0x" << base::hex << (uint64_t)current;
            current = current->next;
            i++;
        }
    }

};