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

#include <ex/pool.h>
#include <ke/bugcheck.h>
#include <hal/kpcr.h>
#include <mm/vm.h>
#include <ptapi/status.h>
#include <string.h>

/*
 * Used for storing in the header so that memory
 * can be freed.
 *
 * @BitBase: Allocation bit base
 * @Gran:    Granularity of page bitmap
 * @Page:    Page we reside in
 */
typedef struct {
    USHORT BitBase : 6;
    USHORT Gran;
    POOL_PAGE *Page;
} PAGE_INFO;

/*
 * Compute the segment level index from the number of
 * bytes requested to be allocated. We use the following
 * formula to do so:
 *
 * ((log2(count) - POOL_MIN_LOG2) - 1) % POOL_LEVEL_MAX
 */
static UCHAR
PoolComputeLevelIndex(USIZE ByteCount)
{
    USIZE TmpCount;
    USIZE Log2Sum = 0;

    /* Compute the base 2 logarithm */
    TmpCount = ByteCount;
    while (TmpCount != 0) {
       TmpCount >>= 1;
       ++Log2Sum;
    }

    /* Subtract the right side */
    TmpCount -= POOL_MIN_LOG2 + 1;
    TmpCount %= POOL_LEVEL_MAX;
    return TmpCount;
}

/*
 * Obtain a reference to the current pool region
 */
static POOL_REGION *
PoolGetCurrentRegion(VOID)
{
    KPCR *Kpcr;

    Kpcr = HalKpcrCurrent();
    return &Kpcr->PoolRegion;
}

/*
 * Allocate a page descriptor
 *
 * @GranBytes:  Granularity in bytes of level
 */
static POOL_PAGE *
PoolAllocatePage(USIZE GranBytes)
{
    PT_STATUS Status;
    POOL_PAGE *PageDesc;
    MM_RANGE Range;
    ULONG PageSizeBitCount;
    ULONG BitmapReserveMask;

    Status = MmAllocatePages(1, &Range);
    if (Status != STATUS_SUCCESS) {
        return NULL;
    }

    PageDesc = RANGE_PTR(&Range);
    PageSizeBitCount = sizeof(*PageDesc) * 8;
    PageSizeBitCount /= GranBytes;
    BitmapReserveMask = (1 << PageSizeBitCount) - 1;
    PageDesc->Bitmap |= BitmapReserveMask;
    return PageDesc;
}

/*
 * Initialize a pool segment
 *
 * @Segment: Pool segment to initialize
 * @Level:   Segment level
 */
static PT_STATUS
PoolSegmentInit(POOL_SEGMENT *Segment, UCHAR Level)
{
    USIZE Gran;

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

    Gran = POOL_LEVEL_GRAN(Level);
    Segment->PageCount = 1;
    Segment->FirstPage = PoolAllocatePage(Gran);
    return STATUS_SUCCESS;
}

/*
 * Allocate memory from a page descriptor
 *
 * @Gran:       Granularity in bytes of level
 * @Page:       Page to allocate from
 * @ByteCount:  Number of bytes to allocate
 * @Info:       Page information result
 */
static VOID *
PoolAllocateFromPage(USIZE Gran, POOL_PAGE *Page, USIZE ByteCount, PAGE_INFO *Info)
{
    USIZE ChunksNeeded, ChunksFound;
    USIZE SetIdx, Idx;
    SSIZE IdxBase = -1;

    if (Info == NULL || Page == NULL) {
        return NULL;
    }

    ChunksNeeded = ByteCount / Gran;

    /* Scan the page for free chunks */
    for (Idx = 0; Idx < sizeof(Page->Bitmap) * 8; ++Idx) {
       if (TESTBIT(&Page->Bitmap, Idx)) {
           ChunksFound = 0;
           IdxBase = -1;
           continue;
       }

       if (IdxBase < 0) {
           IdxBase = Idx;
       }

       ++ChunksFound;
       if (ChunksFound >= ChunksNeeded) {
           break;
       }
    }

    if (ChunksFound < ChunksNeeded) {
        return NULL;
    }

    /* Mark all as allocated */
    for (SetIdx = IdxBase; SetIdx < Idx + 1; ++SetIdx) {
        SETBIT(&Page->Bitmap, SetIdx);
    }

    Info->BitBase = IdxBase;
    Info->Page = Page;
    Info->Gran = Gran;
    return PTR_OFFSET(Page, IdxBase * Gran);
}

/*
 * Allocate memory from a pool segment
 *
 * @ByteCount:  Number of bytes to allocate
 * @LevelIndex: Segment level index
 * @Segment:    Segment to allocate from
 * @Info:       Page information
 */
static VOID *
PoolAllocateFromSegment(USIZE ByteCount, UCHAR LevelIndex, POOL_SEGMENT *Segment, PAGE_INFO *Info)
{
    USIZE ChunkCount, Idx;
    ULONG Gran;
    POOL_PAGE *Page;
    VOID *Ptr;

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

    Gran = POOL_LEVEL_GRAN(LevelIndex);
    Page = Segment->FirstPage;

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

    while (Page != NULL) {
        Ptr = PoolAllocateFromPage(Gran, Page, ByteCount, Info);
        if (Ptr != NULL) {
            return Ptr;
        }

        if (Page->Next == NULL) {
            Page->Next = PoolAllocatePage(Gran);
            if (Page->Next == NULL)
                return NULL;
        }

        Page = Page->Next;
    }

    return NULL;
}

/*
 * TODO: We need per-region locking
 */
VOID *
ExAllocatePoolWithTag(POOL_TYPE Type, USIZE ByteCount, ULONG Tag)
{
    UCHAR SegmentIndex;
    POOL_REGION *Region;
    POOL_SEGMENT *Segment;
    POOL_HEADER *Header;
    PAGE_INFO PageInfo;
    VOID *Base;

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

    ByteCount = ALIGN_UP(
        ByteCount,
        POOL_LEVEL_GRAN(0)
    );

    /* Verify that the pool type is valid */
    switch (Type) {
    case POOL_NON_PAGED:
        break;
    default:
        return NULL;
    }

    Region = PoolGetCurrentRegion();
    SegmentIndex = PoolComputeLevelIndex(ByteCount);
    Segment = &Region->Level[SegmentIndex];
    Base = PoolAllocateFromSegment(
        ByteCount + ALIGN_UP(sizeof(POOL_HEADER), POOL_LEVEL_GRAN(0)),
        SegmentIndex,
        Segment,
        &PageInfo
    );

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

    Header = (POOL_HEADER *)Base;
    Header->Tag = Tag;
    Header->SizeBytes = ByteCount;
    Header->Gran = PageInfo.Gran;
    Header->BitBase = PageInfo.BitBase;
    Header->Page = PageInfo.Page;

    Base = PTR_OFFSET(Base, sizeof(POOL_HEADER));
    RtlMemSet(Base, 0, ByteCount);
    return Base;
}

VOID
ExFreePoolWithTag(VOID *Ptr, ULONG Tag)
{
    POOL_HEADER *Header;
    POOL_PAGE *Page;
    USHORT BitOff, BitMax;

    if (Ptr == NULL) {
        return;
    }

    Header = PTR_NOFFSET(Ptr, sizeof(POOL_HEADER));
    if (Header->Tag != Tag) {
        KeBugCheck(
            BUGCHECK_MISC,
            "tag does not match pool header\n"
        );
    }

    Page = Header->Page;
    if (Page == NULL) {
        KeBugCheck(
            BUGCHECK_MISC,
            "attempted free on malformed header\n"
        );
    }

    BitMax = Header->SizeBytes / Header->Gran;
    BitMax *= 8;

    for (BitOff = Header->BitBase; BitOff < BitMax; ++BitOff) {
        CLRBIT(&Page->Bitmap, BitOff);
    }
}

VOID
ExPoolRegionInit(POOL_REGION *Region)
{
    UCHAR Level;
    POOL_SEGMENT *Segment;
    PT_STATUS Status;

    if (Region == NULL) {
        KeBugCheck(
            BUGCHECK_UNBOUND_RESRC,
            "unbound region in ExPoolRegionInit()\n"
        );
    }

    /* Initialize each segment level */
    for (Level = 0; Level < POOL_LEVEL_MAX; ++Level) {
        Segment = &Region->Level[Level];
        Status = PoolSegmentInit(Segment, Level);

        if (PT_ERROR(Status)) {
            KeBugCheck(
                BUGCHECK_MISC,
                "failed to initialize segment level %d\n",
                Level
            );
        }
    }
}