/*
 * Copyright (c) 2026, Chloe M.
 * Provided under the BSD-3 clause.
 *
 * Description: Pool allocator
 * Author:      Chloe M.
 */

#include <ex/pool.h>
#include <ex/trace.h>
#include <hal/kpcr.h>
#include <mm/vmm.h>
#include <string.h>
#include <bitops.h>

#define DTRACE(Fmt, ...) \
    TRACE("[ POOL ]: " Fmt, ##__VA_ARGS__)

/* Globals */
static USIZE PoolCount = 0;

/*
 * Initialize a pool block
 *
 * @Block: Pool block to initialize
 */
static ST_STATUS
PoolInitBlock(MEMORY_BLOCK *Block)
{
    if (Block == NULL) {
        return STATUS_INVALID_PARAM;
    }

    Block->FirstPage = MmAllocPages(1);
    if (Block->FirstPage == NULL) {
        return STATUS_NO_MEMORY;
    }

    Block->LastPage = Block->FirstPage;
    Block->PageCount = 1;
    RtlMemSet(Block->FirstPage, 0, PAGESIZE);
    return STATUS_SUCCESS;
}

/*
 * Allocates a new memory page to a block and returns it
 *
 * @Block: Block to allocate page to
 */
static MEMORY_PAGE *
BlockNewPage(MEMORY_BLOCK *Block)
{
    MEMORY_PAGE *Page, *Last;

    if (Block == NULL) {
        return NULL;
    }

    Last = Block->LastPage;
    if (Last == NULL) {
        return NULL;
    }

    Page = MmAllocPages(1);
    if (Page == NULL) {
        return NULL;
    }

    RtlMemSet(Page, 0, sizeof(*Page));
    Last->Next = Page;
    Block->LastPage = Page;
    return Page;
}

/*
 * Internally used by AllocatePoolWithTag()
 */
static VOID *
AllocateFromBlock(UCHAR BlockIndex, MEMORY_BLOCK *Block, USIZE Count)
{
    USIZE Gran, Idx;
    USIZE BitsNeeded, BitsFound = 0;
    SSIZE BitBase = -1;
    USIZE BitsSet, Occupied;
    USIZE MaxOccupied;
    MEMORY_PAGE *Page;
    ST_STATUS Status;

    if (Block == NULL || Count == 0) {
        return NULL;
    }

    Gran = LEVEL_GRAN(BlockIndex);
    BitsNeeded = Count / Gran;
    Page = Block->LastPage;

    for (Idx = 0; Idx < sizeof(Page->Bitmap) * 8; ++Idx) {
        if (!TESTBIT(&Page->Bitmap, Idx) && BitsFound < BitsNeeded) {
            if (BitBase < 0)
                BitBase = Idx;

            ++BitsFound;
        } else if (TESTBIT(&Page->Bitmap, Idx)) {
            BitBase = -1;
            BitsFound = 0;
        }

        if (BitsFound >= BitsNeeded) {
            break;
        }
    }

    for (Idx = BitBase; Idx < BitBase + BitsNeeded; ++Idx) {
        SETBIT(&Page->Bitmap, Idx);
    }

    BitsSet = PopCnt(Page->Bitmap);
    Occupied = BitsSet * Gran;
    MaxOccupied = PAGESIZE - sizeof(MEMORY_PAGE);
    MaxOccupied = ALIGN_DOWN(MaxOccupied, Gran);

    /*
     * If there are no more bytes to grab in this page
     */
    if (Occupied >= MaxOccupied || BitsFound < BitsNeeded) {
        Page = BlockNewPage(Block);
        if (Page == NULL) {
            return NULL;
        }

        /* Need to allocate again? */
        if (BitsFound < BitsNeeded) {
            return AllocateFromBlock(BlockIndex, Block, Count);
        }
    }

    return (BitsFound >= BitsNeeded)
        ? PTR_OFFSET(Page, sizeof(MEMORY_PAGE) + (BitBase * Gran))
        : NULL;
}

/*
 * Internally used by ExAllocatePoolWithTag()
 */
static VOID *
AllocatePoolWithTag(MEMORY_POOL *Pool, POOL_TYPE Type, USIZE ByteCount, ULONG Tag)
{
    USIZE TmpCount;
    UCHAR BlockIndex = 0;
    MEMORY_BLOCK *Block;

    if (Pool == NULL || ByteCount == 0) {
        return NULL;
    }

    /*
     * Obtain the block index with the formula:
     *
     * i = (log2(count) - POOL_MIN_LOG2) - 1
     */
    ByteCount = ALIGN_UP(ByteCount, LEVEL_GRAN(0));
    TmpCount = ByteCount;
    while (TmpCount != 0) {
        ++BlockIndex;
        TmpCount >>= 1;
    }

    BlockIndex -= (POOL_MIN_LOG2 + 1);
    BlockIndex %= LEVEL_COUNT;
    Block = &Pool->BlockLevels[BlockIndex];
    return AllocateFromBlock(BlockIndex, Block, ByteCount);
}

VOID *
ExAllocatePoolWithTag(POOL_TYPE Type, USIZE ByteCount, ULONG Tag)
{
    KPCR *ThisCore;
    MEMORY_POOL *AllocPool;
    MEMORY_HEADER *Header;
    VOID *Base;

    if (ByteCount == 0) {
        return NULL;
    }

    /* Verify that the pool type is valid */
    switch (Type) {
    case POOL_NON_PAGED:
        break;
    default:
        DTRACE("got bad pool type in allocation\n");
        return NULL;
    }

    ThisCore = HalKpcrCurrent();
    AllocPool = &ThisCore->AllocPool;
    KeSpinLockAcquire(&AllocPool->Lock, false);

    /*
     * Ensure that the byte count is aligned by the minimum granularity
     * that we can allocate. Furthermore, ensure that it does not exceed
     * the size of a single page.
     */
    ByteCount = ALIGN_UP(ByteCount + sizeof(*Header), LEVEL_GRAN(0));
    if (ByteCount >= PAGESIZE) {
        DTRACE(
            "allocation in pool #%d exceeds page size\n",
            AllocPool->Id
        );

        KeSpinLockRelease(&AllocPool->Lock);
        return NULL;
    }

    Base = AllocatePoolWithTag(AllocPool, Type, ByteCount, Tag);
    KeSpinLockRelease(&AllocPool->Lock);

    if (Base == NULL) {
        return NULL;
    }

    Header = (MEMORY_HEADER *)Base;
    Header->Magic = POOL_MAGIC;
    Header->MemoryTag = Tag;
    Header->Length = ByteCount;
    return PTR_OFFSET(Base, sizeof(MEMORY_HEADER));
}

ST_STATUS
ExPoolInit(MEMORY_POOL *Pool)
{
    UCHAR LevelIdx;
    MEMORY_BLOCK *Block;
    ST_STATUS Status;

    if (Pool == NULL) {
        return STATUS_INVALID_PARAM;
    }

    Pool->Id = PoolCount;
    KeSpinLockInit(&Pool->Lock, "ex-pool");
    DTRACE("bringing up pool #%d\n", Pool->Id);

    /* Initialize the pools */
    for (LevelIdx = 0; LevelIdx < LEVEL_COUNT; ++LevelIdx) {
        Block = &Pool->BlockLevels[LevelIdx];
        Status = PoolInitBlock(Block);

        /* XXX: We may wanna clean up our mess we made */
        if (Status != STATUS_SUCCESS) {
            DTRACE("critical: failed to initialize pool\n");
            return STATUS_NO_MEMORY;
        }
    }

    DTRACE("brought up %d levels for pool #%d : ok\n", LEVEL_COUNT, Pool->Id);
    ++PoolCount;
    return STATUS_SUCCESS;
}