/*
 * PROJECT: Alcyone System Kernel
 * LICENSE: BSD Clause 3
 * PURPOSE: Cache Controllerm, Cache  Utility Handler.
 * NT KERNEL: 5.11.9360
 * COPYRIGHT:  2023-2029 Dibymartanda Samanta <>
 */ 

class BitmapRangeManager {
private:
static constexpr LONGLONG CalculateBasePage(LONGLONG InputPage) {
        return (InputPage & ~(0x1000 - 1));
    }
public:
    static PBITMAP_RANGE FindBitmapRangeToDirty(PMBCB Mbcb, LONGLONG InputPage, PULONG* pBitmap) {
        PBITMAP_RANGE CurrentRange = nullptr;
        PBITMAP_RANGE NewRange = nullptr;
        PLIST_ENTRY HeadList = nullptr;
        LONGLONG BasePage = {0};


        HeadList = &Mbcb->BitmapRanges;
        BasePage = CalculateBasePage(InputPage);

        CurrentRange = CONTAINING_RECORD(Mbcb->BitmapRanges.Flink, BITMAP_RANGE, Links);

        while(true){
            if (BasePage == CurrentRange->BasePage)
                return CurrentRange;

            if (CurrentRange->DirtyPages || NewRange) {
                if (BasePage > CurrentRange->BasePage)
                    HeadList = &CurrentRange->Links;
            }
            else {
                NewRange = CurrentRange;
            }

            if (CurrentRange->Links.Flink == &Mbcb->BitmapRanges)
                break;

            CurrentRange = CONTAINING_RECORD(CurrentRange->Links.Flink, BITMAP_RANGE, Links);

            if (CurrentRange->BasePage > BasePage && NewRange)
                break;
        } 

        if (NewRange) {
            RemoveEntryList(&NewRange->Links);
        }
        else {
            NewRange = reinterpret_cast<PBITMAP_RANGE>(ExAllocatePoolWithTag(NonPagedPool, sizeof(*NewRange), 'rBcC'));
            if (!NewRange) {
                return nullptr;
            }

            RtlZeroMemory(NewRange, sizeof(*NewRange));
        }

        InsertHeadList(HeadList, &NewRange->Links);

        NewRange->BasePage = BasePage;
        NewRange->FirstDirtyPage = 0xFFFFFFFF;
        NewRange->LastDirtyPage = 0;

        if (!NewRange->Bitmap) {

            NewRange->Bitmap = *pBitmap;
            *pBitmap = nullptr;
        }

        return NewRange;
    }
};

/*  Internal Functions */
LONG
CcCopyReadExceptionFilter(
    _In_ PEXCEPTION_POINTERS ExceptionInfo,
    _Out_ NTSTATUS* OutStatus)
{
   

    NT_ASSERT(!NT_SUCCESS(*OutStatus));

    return EXCEPTION_EXECUTE_HANDLER;
}

PBITMAP_RANGE
NTAPI
CcFindBitmapRangeToDirty(
    _In_ PMBCB Mbcb,
    _In_ LONGLONG InputPage,
    _Inout_ PULONG* PBITMAP)
{
    return BitmapRangeManager::FindBitmapRangeToDirty(Mbcb, InputPage, pBitmap);
    
}

VOID
NTAPI
CcSetDirtyInMask(
    _In_ PSHARED_CACHE_MAP SharedCacheMap,
    _In_ PLARGE_INTEGER FileOffset,
    _In_ ULONG Length)
{
    KLOCK_QUEUE_HANDLE LockHandle = {0};
    LARGE_INTEGER EndOffset = {0};
    ULONGLONG StartPage ={0}; 
    ULONGLONG EndPage = {0};
    ULONG CurrentPage = {0};
    PMBCB Mbcb = nullptr;
    PBITMAP_RANGE BitmapRange = nullptr;
    PULONG Bitmap = nullptr;
    PULONG VacbLevel = nullptr;
    ULONG BitMask = {0};



    // Calculate start and end pages
    StartPage = FileOffset->QuadPart >> PAGE_SHIFT;
    EndOffset.QuadPart = FileOffset->QuadPart + Length;
    EndPage = (EndOffset.QuadPart - 1) >> PAGE_SHIFT;

    for(;;)
    {
        if (SharedCacheMap->SectionSize.QuadPart > VACB_MAPPING_GRANULARITY)
        {
            if (!CcPrefillVacbLevelZone(1, &LockHandle, 0, 0, 0))
                return;
            VacbLevel = CcAllocateVacbLevel(0);
            KeLowerIrql(LockHandle.OldIrql);
        }

        KeAcquireInStackQueuedSpinLock(&SharedCacheMap->BcbSpinLock, &LockHandle);

        Mbcb = SharedCacheMap->Mbcb;
        if (Mbcb == nullptr)
        {
            Mbcb = CcAllocateInitializeBcb();
            if (Mbcb == nullptr)
                goto ReleaseAndExit;

            Mbcb->NodeTypeCode = CACHE_NTC_BCB;
            InitializeListHead(&Mbcb->BitmapRanges);
            InsertHeadList(&Mbcb->BitmapRanges, &Mbcb->BitmapRange1.Links);
            Mbcb->BitmapRange1.FirstDirtyPage = (ULONG)-1;
            Mbcb->BitmapRange1.Bitmap = (PULONG)&Mbcb->BitmapRange2;
            SharedCacheMap->Mbcb = Mbcb;
        }

        if (EndPage < 512 || Mbcb->NodeTypeCode == 0x02F9)
        {
            BitmapRange = CcFindBitmapRangeToDirty(Mbcb, StartPage, &Bitmap);
            if (BitmapRange == nullptr)
                break;

            if (StartPage < BitmapRange->BasePage + BitmapRange->FirstDirtyPage)
                BitmapRange->FirstDirtyPage = StartPage - BitmapRange->BasePage;
            if (EndPage > BitmapRange->BasePage + BitmapRange->LastDirtyPage)
                BitmapRange->LastDirtyPage = EndPage - BitmapRange->BasePage;

            if (SharedCacheMap->DirtyPages == 0)
            {
                CcScheduleLazyWriteScan(FALSE);
                RemoveEntryList(&SharedCacheMap->SharedCacheMapLinks);
                InsertTailList(&CcDirtySharedCacheMapList, &SharedCacheMap->SharedCacheMapLinks);
                Mbcb->ResumeWritePage = StartPage;
            }

            Bitmap = &BitmapRange->Bitmap[(StartPage - BitmapRange->BasePage) >> 5];
            BitMask = 1 << (StartPage & 0x1F);

            for (CurrentPage = StartPage; CurrentPage <= EndPage; CurrentPage++)
            {
                if ((*Bitmap & BitMask) == 0)
                {
                    CcTotalDirtyPages++;
                    Mbcb->DirtyPages++;
                    BitmapRange->DirtyPages++;
                    SharedCacheMap->DirtyPages++;
                    *Bitmap |= BitMask;
                }
                BitMask <<= 1;
                if (BitMask == 0)
                {
                    Bitmap++;
                    BitMask = 1;
                }
            }
        }
        else
        {
            // Handle large files (>2MB)
            if (Mbcb->BitmapRange1.DirtyPages)
            {
                RtlCopyMemory(VacbLevel, Mbcb->BitmapRange1.Bitmap,  (2 * sizeof(BITMAP_RANGE)));
                RtlZeroMemory(Mbcb->BitmapRange1.Bitmap,  (2 * sizeof(BITMAP_RANGE)));
            }
            Mbcb->BitmapRange1.Bitmap = VacbLevel;

            // Initialize BitmapRange2
            InsertTailList(&Mbcb->BitmapRanges, &Mbcb->BitmapRange2.Links);
            Mbcb->BitmapRange2.BasePage = (ULONGLONG)-1;
            Mbcb->BitmapRange2.FirstDirtyPage = (ULONG)-1;

            // Initialize BitmapRange3
            InsertTailList(&Mbcb->BitmapRanges, &Mbcb->BitmapRange3.Links);
            Mbcb->BitmapRange3.BasePage = (ULONGLONG)-1;
            Mbcb->BitmapRange3.FirstDirtyPage = (ULONG)-1;

            VacbLevel = nullptr;
            Mbcb->NodeTypeCode = 0x02F9;

            KeReleaseInStackQueuedSpinLock(&LockHandle);
            continue;
        }

        // Update ValidDataGoal if necessary
        if (EndOffset.QuadPart > SharedCacheMap->ValidDataGoal.QuadPart)
        {
            SharedCacheMap->ValidDataGoal = EndOffset;
        }

        break;
    } 

ReleaseAndExit:
    if (VacbLevel != nullptr)
    {
        *VacbLevel = (ULONG)CcVacbLevelFreeList;
        CcVacbLevelEntries++;
        CcVacbLevelFreeList = VacbLevel;
    }
    KeReleaseInStackQueuedSpinLock(&LockHandle);
}

BOOLEAN
NTAPI
CcMapAndCopy(
    _In_ PSHARED_CACHE_MAP SharedCacheMap,
    _In_ PVOID Buffer,
    _In_ PLARGE_INTEGER FileOffset,
    _In_ ULONG Length,
    _In_ ULONG Flags,
    _In_ PFILE_OBJECT FileObject,
    _In_ PLARGE_INTEGER ValidDataLength,
    _In_ BOOLEAN Wait)
{
   NT_DBGBREAK("UNIMPLEMENTED\n");
}

 
         
           

/* EXTERNAL API FUNCTIONS */

BOOLEAN
NTAPI
CcCanIWrite(
    _In_ PFILE_OBJECT FileObject,
    _In_ ULONG BytesToWrite,
    _In_ BOOLEAN Wait,
    _In_ UCHAR Retrying)
{
    PFSRTL_COMMON_FCB_HEADER FsContext = nullptr; 
    PSHARED_CACHE_MAP SharedCacheMap = nullptr;
    DEFERRED_WRITE DeferredWrite = {0};
    KEVENT Event = {0};
    ULONG WriteSize = {0};
    ULONG Pages = {0};
    KIRQL OldIrql = {0};
    BOOLEAN IsSmallThreshold = false;


    /* Quick checks for immediate return */
    if (FileObject->Flags & FO_WRITE_THROUGH)
        return true;

    if (IoIsFileOriginRemote(FileObject) && Retrying < 0xFD)
        return true;

    /* Calculate size and pages */
    WriteSize = min(BytesToWrite, 0x40000);
    Pages = ROUND_UP(WriteSize, PAGE_SIZE) / PAGE_SIZE;

    FsContext = FileObject->FsContext;

    /* Check threshold conditions */
    if (Retrying >= 0xFE || (FsContext->Flags & FSRTL_FLAG_LIMIT_MODIFIED_PAGES))
    {
        if (Retrying != 0xFF)
        {
            OldIrql = KeAcquireQueuedSpinLock(LockQueueMasterLock);
        }

        if (FileObject->SectionObjectPointer)
        {
            SharedCacheMap = FileObject->SectionObjectPointer->SharedCacheMap;
            if (SharedCacheMap && 
                SharedCacheMap->DirtyPageThreshold && 
                SharedCacheMap->DirtyPages)
            {
                if (SharedCacheMap->DirtyPageThreshold < (SharedCacheMap->DirtyPages + Pages))
                {
                    IsSmallThreshold = true;
                }
            }
        }

        if (Retrying != 0xFF)
        {
            KeReleaseQueuedSpinLock(LockQueueMasterLock, OldIrql);
        }
    }

    /* Check if we can write immediately */
    if ((Retrying || IsListEmpty(&CcDeferredWrites)) &&
        (CcDirtyPageThreshold > (CcTotalDirtyPages + Pages)))
    {
        if ((MmAvailablePages > MmThrottleTop && !IsSmallThreshold) ||
            (MmModifiedPageListHead.Total < 1000 &&
             MmAvailablePages > MmThrottleBottom &&
             !IsSmallThreshold))
        {
            return true;
        }
    }

    if (!Wait)
        return FALSE;

    /* Initialize deferred write */
    if (IsListEmpty(&CcDeferredWrites))
    {
        OldIrql = KeAcquireQueuedSpinLock(LockQueueMasterLock);
        CcScheduleLazyWriteScan(TRUE);
        KeReleaseQueuedSpinLock(LockQueueMasterLock, OldIrql);
    }

    KeInitializeEvent(&Event, NotificationEvent, false);

    DeferredWrite.NodeTypeCode = NODE_TYPE_DEFERRED_WRITE;
    DeferredWrite.NodeByteSize = sizeof(DEFERRED_WRITE);
    DeferredWrite.FileObject = FileObject;
    DeferredWrite.BytesToWrite = BytesToWrite;
    DeferredWrite.Event = &Event;
    DeferredWrite.LimitModifiedPages = (FsContext->Flags & FSRTL_FLAG_LIMIT_MODIFIED_PAGES) != 0;

    /* Insert into deferred writes list */
    /* Insert into deferred writes list */
    if (Retrying)
    {
        ExInterlockedInsertHeadList(&CcDeferredWrites,
                                  &DeferredWrite.DeferredWriteLinks,
                                  &CcDeferredWriteSpinLock);
    }
    else
    {
        ExInterlockedInsertTailList(&CcDeferredWrites,
                                  &DeferredWrite.DeferredWriteLinks,
                                  &CcDeferredWriteSpinLock);
    }

    /* Wait for the write to complete */
    do
    {
        CcPostDeferredWrites();
    } while (KeWaitForSingleObject(&Event,
                                 Executive,
                                 KernelMode,
                                 false,
                                 &CcIdleDelay) != STATUS_SUCCESS);

    return TRUE;
}


BOOLEAN
NTAPI
CcCopyRead(
    _In_ PFILE_OBJECT FileObject,
    _In_ PLARGE_INTEGER FileOffset,
    _In_ ULONG Length,
    _In_ BOOLEAN Wait,
    _Out_ PVOID Buffer,
    _Out_ IO_STATUS_BLOCK* OutIoStatus)
{
               NT_DBGBREAK("UNIMPLEMENTED\n");
}


BOOLEAN
NTAPI
CcCopyWrite(
    _In_ PFILE_OBJECT FileObject,
    _In_ PLARGE_INTEGER FileOffset,
    _In_ ULONG Length,
    _In_ BOOLEAN Wait,
    _In_ PVOID Buffer)
{
    NT_DBGBREAK("UNIMPLEMENTED\n");
}


VOID
CcDeferWrite(
    _In_ PFILE_OBJECT FileObject,
    _In_ PCC_POST_DEFERRED_WRITE PostRoutine,
    _In_ PVOID Context1,
    _In_ PVOID Context2,
    _In_ ULONG BytesToWrite,
    _In_ BOOLEAN Retrying
    )
{
    
    KIRQL OldIrql;


    PDEFERRED_WRITE  DeferredWrite = ExAllocatePoolWithTag(NonPagedPoolNx, sizeof(DEFERRED_WRITE), 'wDcC');
    if (DeferredWrite != nullptr)
    {
        RtlZeroMemory(DeferredWrite, sizeof(DEFERRED_WRITE));

        DeferredWrite->NodeTypeCode = NODE_TYPE_DEFERRED_WRITE;
        DeferredWrite->NodeByteSize = sizeof(DEFERRED_WRITE);
        DeferredWrite->FileObject = FileObject;
        DeferredWrite->BytesToWrite = BytesToWrite;
        DeferredWrite->PostRoutine = PostRoutine;
        DeferredWrite->Context1 = Context1;
        DeferredWrite->Context2 = Context2;

        if (Retrying)
            ExInterlockedInsertHeadList(&CcDeferredWrites, 
                                        &DeferredWrite->DeferredWriteLinks, 
                                        &CcDeferredWriteSpinLock);
        else
            ExInterlockedInsertTailList(&CcDeferredWrites, 
                                        &DeferredWrite->DeferredWriteLinks, 
                                        &CcDeferredWriteSpinLock);

        CcPostDeferredWrites();

        OldIrql = KeAcquireQueuedSpinLock(LockQueueMasterLock);
        CcScheduleLazyWriteScan(TRUE);
        KeReleaseQueuedSpinLock(LockQueueMasterLock, OldIrql);
    }
    else
    {
        PostRoutine(Context1, Context2);
    }
}

VOID
NTAPI
CcFastCopyRead(
    _In_ PFILE_OBJECT FileObject,
    _In_ ULONG FileOffset,
    _In_ ULONG Length,
    _In_ ULONG PageCount,
    _Out_ PVOID Buffer,
    _Out_ PIO_STATUS_BLOCK IoStatus)
{
    NT_DBGBREAK("UNIMPLEMENTED\n");
}


VOID
NTAPI
CcFastCopyWrite(
    _In_ PFILE_OBJECT FileObject,
    _In_ ULONG FileOffset,
    _In_ ULONG Length,
    _In_ PVOID InBuffer)
{
    NT_DBGBREAK("UNIMPLEMENTED\n");
}