// FAT allocation, cache, and metadata core shared by the filesystem operation translation units.

#include <stdbool.h>
#include "fat_internal.h"

// FAT sector cache (fixes #1 and #5)
// 32 direct-mapped slots, indexed by (lba % FAT_CACHE_SLOTS).
// Write-back: FatWriteEntry modifies the in-cache bytes and marks the slot
// dirty. FatFlushCache() propagates dirty slots to every FAT copy at once,
// replacing the original per-entry 2x(read+write) pattern with a single
// write per dirty sector per FAT copy.

#define FAT_CACHE_SLOTS 32

typedef struct {
    FAT_VOL* vol;
    uint64_t lba;
    uint8_t  data[512];
    bool     valid;
    bool     dirty;
} FAT_CACHE_ENTRY;

static FAT_CACHE_ENTRY gFatCache[FAT_CACHE_SLOTS];

// The FAT sector cache is global and direct-mapped. mmap-backed video can fault
// file pages from multiple Bink worker threads at the same time; without
// serialization one thread can evict/reload a cache slot while another thread is
// still reading a FAT entry from that slot, corrupting cluster-chain traversal.
// Use a sleepable mutex rather than a spinlock because cache misses perform disk
// I/O. The static mutex avoids allocating an object from the paging path.
static MUTEX    gFatCacheMutex;
static SPINLOCK gFatCacheMutexInitLock = SPINLOCK_INIT;
static int      gFatCacheMutexReady;

static void FatCacheEnsureMutex(void) {
    if (gFatCacheMutexReady)
        return;

    uintptr_t flags = KeAcquireSpinlock(&gFatCacheMutexInitLock);
    if (!gFatCacheMutexReady) {
        memset(&gFatCacheMutex, 0, sizeof(gFatCacheMutex));
        ObInitializeObjectHeader(&gFatCacheMutex.Header, ObjectTypeMutex);
        gFatCacheMutex.Lock = (SPINLOCK)SPINLOCK_INIT;
        gFatCacheMutex.Owner = (struct THREAD*)0;
        LIST_INIT(&gFatCacheMutex.WaitList);
        LIST_INIT(&gFatCacheMutex.OwnerLink);
        gFatCacheMutex.DonationClass = KeSchedulingNormal;
        gFatCacheMutex.DonationPriority = PriorityLow;
        gFatCacheMutex.HasDonation = 0;
        gFatCacheMutex.Abandoned = 0;
        gFatCacheMutexReady = 1;
    }
    KeReleaseSpinlock(&gFatCacheMutexInitLock, flags);
}

static int FatCacheAcquire(void) {
    if (!KeCanWait())
        return 0;
    FatCacheEnsureMutex();
    return KE_MUTEX_ACQUIRE_SUCCEEDED(KeAcquireMutex(&gFatCacheMutex, KE_WAIT_INFINITE));
}

static void FatCacheRelease(int held) {
    if (held)
        KeReleaseMutex(&gFatCacheMutex);
}

// Raw (non-cached) sector I/O — forward declarations.
static KSTATUS ReadSectorsRaw (FAT_VOL* v, uint64_t lba, uint32_t count, void* buf);
static KSTATUS WriteSectorsRaw(FAT_VOL* v, uint64_t lba, uint32_t count, const void* buf);

// Read one FAT sector through the cache.
// Caller must hold the FAT cache mutex while using the returned pointer.
static KSTATUS FatCacheRead(FAT_VOL* v, uint64_t lba, uint8_t** out) {
    uint32_t        slot = (uint32_t)(lba % FAT_CACHE_SLOTS);
    FAT_CACHE_ENTRY* e   = &gFatCache[slot];

    if (e->valid && e->vol == v && e->lba == lba) {
        *out = e->data;
        return NTLSTATUS_SUCCESS;
    }

    // Evict current occupant: single-copy writeback on collision.
    // The full multi-copy flush happens in FatFlushCache().
    if (e->valid && e->dirty && e->vol) {
        WriteSectorsRaw(e->vol, e->lba, 1, e->data);
        e->dirty = false;
    }

    KSTATUS s = ReadSectorsRaw(v, lba, 1, e->data);
    if (!KSUCCESS(s)) { e->valid = false; e->vol = (FAT_VOL*)0; return s; }

    e->vol   = v;
    e->lba   = lba;
    e->valid = true;
    e->dirty = false;
    *out = e->data;
    return NTLSTATUS_SUCCESS;
}

static void FatCacheMarkDirty(FAT_VOL* v, uint64_t lba) {
    uint32_t slot = (uint32_t)(lba % FAT_CACHE_SLOTS);
    if (gFatCache[slot].valid && gFatCache[slot].vol == v && gFatCache[slot].lba == lba)
        gFatCache[slot].dirty = true;
}

// Invalidates every FAT sector-cache entry.
void FatCacheInvalidate(void) {
    int held = FatCacheAcquire();
    for (int i = 0; i < FAT_CACHE_SLOTS; i++) {
        gFatCache[i].vol = (FAT_VOL*)0;
        gFatCache[i].valid = false;
        gFatCache[i].dirty = false;
    }
    FatCacheRelease(held);
}

// Discards cache entries owned by one dead or destroyed mounted FAT instance.
void FatCacheInvalidateVolume(FAT_VOL* volume) {
    if (!volume)
        return;
    int held = FatCacheAcquire();
    for (int i = 0; i < FAT_CACHE_SLOTS; i++) {
        if (gFatCache[i].vol != volume)
            continue;
        gFatCache[i].vol = (FAT_VOL*)0;
        gFatCache[i].valid = false;
        gFatCache[i].dirty = false;
    }
    FatCacheRelease(held);
}

// FatFlushCache — write every dirty slot to all FAT copies, clear dirty flag.
// Defined after FAT_VOL is fully declared (below).
void FatFlushCache(FAT_VOL* v);

// Raw sector I/O (bypasses FAT cache; used for file data and cache eviction).
static KSTATUS ReadSectorsRaw(FAT_VOL* v, uint64_t lba,
                               uint32_t count, void* buf) {
    if (!v) return NTLSTATUS_INVALID_PARAMETER;
    return IoReadDiskSectors(v->DriveIndex, lba, count, buf);
}

static KSTATUS WriteSectorsRaw(FAT_VOL* v, uint64_t lba,
                                uint32_t count, const void* buf) {
    if (!v) return NTLSTATUS_INVALID_PARAMETER;
    return IoWriteDiskSectors(v->DriveIndex, lba, count, buf);
}

// Convenience aliases for file-data I/O.
KSTATUS FatReadSectors (FAT_VOL* v, uint64_t lba, uint32_t n, void* b) { return ReadSectorsRaw (v, lba, n, b); }
// Updates sectors.
KSTATUS FatWriteSectors(FAT_VOL* v, uint64_t lba, uint32_t n, const void* b) { return WriteSectorsRaw(v, lba, n, b); }

const char* FatVariantName(FAT_VARIANT variant) {
    switch (variant) {
        case FatVariant16: return "FAT16";
        case FatVariant32: return "FAT32";
        case FatVariant12: return "FAT12";
        default: return "FAT";
    }
}

static uint32_t FatEocThreshold(FAT_VOL* v) {
    return v && v->Variant == FatVariant16 ? FAT16_EOC : FAT32_EOC;
}

uint32_t FatEocMark(FAT_VOL* v) {
    return v && v->Variant == FatVariant16 ? FAT16_EOC_MARK : FAT32_EOC_MARK;
}

bool FatIsDataCluster(FAT_VOL* v, uint32_t cluster) {
    return v && cluster >= 2u && cluster < FatEocThreshold(v);
}

uint32_t FatDirentFirstCluster(FAT_VOL* v, const DIRENT* d) {
    uint32_t lo = d ? d->FirstClusterLo : 0u;
    uint32_t hi = (v && v->Variant == FatVariant32 && d) ? d->FirstClusterHi : 0u;
    return (hi << 16) | lo;
}

void FatSetDirentFirstCluster(FAT_VOL* v, DIRENT* d, uint32_t cluster) {
    if (!d) return;
    d->FirstClusterLo = (uint16_t)(cluster & 0xFFFFu);
    d->FirstClusterHi = (v && v->Variant == FatVariant32) ? (uint16_t)(cluster >> 16) : 0u;
}

bool FatIsFixedRoot(FAT_VOL* v, uint32_t dirCluster) {
    return v && v->Variant == FatVariant16 && dirCluster == FAT_FIXED_ROOT_CLUSTER;
}

uint64_t FatClusterToLba(FAT_VOL* v, uint32_t cluster);
uint32_t FatEntry(FAT_VOL* v, uint32_t cluster);

bool FatDirectoryFirstSegment(FAT_VOL* v, uint32_t dirCluster,
                                    uint32_t* cursor, uint64_t* lba,
                                    uint32_t* sectors) {
    if (!v || !cursor || !lba || !sectors) return false;
    if (FatIsFixedRoot(v, dirCluster)) {
        if (!v->RootDirSectors) return false;
        *cursor = 0;
        *lba = v->RootDirStart;
        *sectors = 1;
        return true;
    }
    if (!FatIsDataCluster(v, dirCluster)) return false;
    *cursor = dirCluster;
    *lba = FatClusterToLba(v, dirCluster);
    *sectors = v->SectorsPerCluster;
    return true;
}

bool FatDirectoryNextSegment(FAT_VOL* v, uint32_t dirCluster,
                                   uint32_t* cursor, uint64_t* lba,
                                   uint32_t* sectors) {
    if (!v || !cursor || !lba || !sectors) return false;
    if (FatIsFixedRoot(v, dirCluster)) {
        uint32_t next = *cursor + 1u;
        if (next >= v->RootDirSectors) return false;
        *cursor = next;
        *lba = v->RootDirStart + next;
        *sectors = 1;
        return true;
    }
    uint32_t next = FatEntry(v, *cursor);
    if (!FatIsDataCluster(v, next)) return false;
    *cursor = next;
    *lba = FatClusterToLba(v, next);
    *sectors = v->SectorsPerCluster;
    return true;
}

uint64_t FatClusterToLba(FAT_VOL* v, uint32_t cluster) {
    return v->DataStart + (uint64_t)(cluster - 2) * v->SectorsPerCluster;
}

// FatFlushCache — needs full FAT_VOL definition.
void FatFlushCache(FAT_VOL* v) {
    uint64_t fat0Start = (uint64_t)v->PartitionLba + v->ReservedSectors;
    int held = FatCacheAcquire();

    for (int i = 0; i < FAT_CACHE_SLOTS; i++) {
        FAT_CACHE_ENTRY* e = &gFatCache[i];
        if (!e->valid || !e->dirty || e->vol != v) continue;
        if (e->lba < fat0Start) continue;

        uint64_t sectorIdx = e->lba - fat0Start;
        if (sectorIdx >= v->FatSize) continue;

        for (uint32_t f = 0; f < v->NumFats; f++) {
            uint64_t lba = fat0Start + (uint64_t)f * v->FatSize + sectorIdx;
            WriteSectorsRaw(v, lba, 1, e->data);
        }
        e->dirty = false;
    }

    FatCacheRelease(held);
}

// FAT read / write (through cache).

uint32_t FatEntry(FAT_VOL* v, uint32_t cluster) {
    if (!v || !FatIsDataCluster(v, cluster)) return 0;

    uint32_t bytesPerEntry = v->Variant == FatVariant16 ? 2u : 4u;
    uint32_t fatOffset = cluster * bytesPerEntry;
    uint64_t sector    = (uint64_t)v->PartitionLba + v->ReservedSectors
                       + fatOffset / v->BytesPerSector;
    uint32_t offInSec  = fatOffset % v->BytesPerSector;

    int held = FatCacheAcquire();
    uint8_t* sec;
    KSTATUS s = FatCacheRead(v, sector, &sec);
    uint32_t value = 0;
    if (KSUCCESS(s)) {
        if (v->Variant == FatVariant16)
            value = *(uint16_t*)(sec + offInSec);
        else
            value = *(uint32_t*)(sec + offInSec) & 0x0FFFFFFFu;
    }
    FatCacheRelease(held);
    return value;
}

// Updates fat Entry.
KSTATUS FatWriteEntry(FAT_VOL* v, uint32_t cluster, uint32_t value) {
    if (!v || !FatIsDataCluster(v, cluster)) return NTLSTATUS_INVALID_PARAMETER;

    uint32_t bytesPerEntry = v->Variant == FatVariant16 ? 2u : 4u;
    uint32_t fatOffset = cluster * bytesPerEntry;
    uint64_t sector    = (uint64_t)v->PartitionLba + v->ReservedSectors
                       + fatOffset / v->BytesPerSector;
    uint32_t offInSec  = fatOffset % v->BytesPerSector;

    int held = FatCacheAcquire();
    uint8_t* sec;
    KSTATUS s = FatCacheRead(v, sector, &sec);
    if (!KSUCCESS(s)) {
        FatCacheRelease(held);
        return s;
    }

    if (v->Variant == FatVariant16) {
        *(uint16_t*)(sec + offInSec) = (uint16_t)(value & 0xFFFFu);
    } else {
        uint32_t existing = *(uint32_t*)(sec + offInSec);
        *(uint32_t*)(sec + offInSec) = (existing & 0xF0000000u)
                                      | (value    & 0x0FFFFFFFu);
    }
    FatCacheMarkDirty(v, sector);
    FatCacheRelease(held);
    return NTLSTATUS_SUCCESS;
}

// Cluster allocation (fix #3: next-free hint + two-phase wrap-around scan).
uint32_t FatAllocClusterEx(FAT_VOL* v, bool zeroCluster) {
    if (v->TotalClusters == 0) return 0;

    uint32_t total = v->TotalClusters;
    uint32_t start = v->NextFreeHint;
    if (start < 2 || start >= total + 2) start = 2;

    for (uint32_t pass = 0; pass < 2; pass++) {
        uint32_t lo = (pass == 0) ? start   : 2;
        uint32_t hi = (pass == 0) ? total+2 : start;
        for (uint32_t c = lo; c < hi; c++) {
            if (FatEntry(v, c) == FAT_FREE) {
                if (!KSUCCESS(FatWriteEntry(v, c, FatEocMark(v)))) return 0;
                v->NextFreeHint = c + 1;

                // Normal file growth zeroes newly allocated clusters so old
                // disk contents are never exposed. Swap/pagefile creation is
                // the one exception: every allocated slot is overwritten before
                // it can ever be read back, so skipping these writes avoids a
                // slow boot-time full-pagefile zero pass.
                if (zeroCluster) {
                    uint8_t* zero = (uint8_t*)MmAllocatePool(v->BytesPerCluster);
                    if (zero) {
                        memset(zero, 0, v->BytesPerCluster);
                        FatWriteSectors(v, FatClusterToLba(v, c), v->SectorsPerCluster, zero);
                        MmFreePool(zero);
                    }
                }
                return c;
            }
        }
    }
    return 0;
}

uint32_t FatAllocCluster(FAT_VOL* v) {
    return FatAllocClusterEx(v, 1);
}

uint32_t FatAppendClusterEx(FAT_VOL* v, uint32_t lastCluster, bool zeroCluster) {
    uint32_t newCluster = FatAllocClusterEx(v, zeroCluster);
    if (!newCluster) return 0;
    if (!KSUCCESS(FatWriteEntry(v, lastCluster, newCluster))) return 0;
    return newCluster;
}

// Appends cluster.
uint32_t FatAppendCluster(FAT_VOL* v, uint32_t lastCluster) {
    return FatAppendClusterEx(v, lastCluster, 1);
}

// Returns chain End.
uint32_t FatGetChainEnd(FAT_VOL* v, uint32_t start, uint32_t* outCount) {
    uint32_t cluster = start, prev = start, count = 0;
    while (FatIsDataCluster(v, cluster)) {
        prev = cluster; count++;
        cluster = FatEntry(v, cluster);
    }
    if (outCount) *outCount = count;
    return prev;
}

// Releases cluster Chain.
void FatFreeClusterChain(FAT_VOL* v, uint32_t start) {
    uint32_t cluster = start;
    while (FatIsDataCluster(v, cluster)) {
        uint32_t next = FatEntry(v, cluster);
        FatWriteEntry(v, cluster, FAT_FREE);
        if (cluster < v->NextFreeHint) v->NextFreeHint = cluster;
        cluster = next;
    }
}

// FAT timestamp and DIRENT metadata helpers.

static bool FatIsLeapYear(uint32_t year) {
    return ((year % 4u) == 0u && (year % 100u) != 0u) || ((year % 400u) == 0u);
}

static uint32_t FatDaysInMonth(uint32_t year, uint32_t month) {
    static const uint8_t days[12] = { 31,28,31,30,31,30,31,31,30,31,30,31 };
    if (month < 1u || month > 12u) return 31u;
    if (month == 2u && FatIsLeapYear(year)) return 29u;
    return days[month - 1u];
}

static uint64_t FatCalendarToEpochSeconds(uint32_t year, uint32_t month, uint32_t day,
                                          uint32_t hour, uint32_t minute, uint32_t second) {
    uint64_t days = 0;
    for (uint32_t y = 1970u; y < year; y++)
        days += FatIsLeapYear(y) ? 366u : 365u;
    for (uint32_t m = 1u; m < month; m++)
        days += FatDaysInMonth(year, m);
    days += day - 1u;
    return (((days * 24u) + hour) * 60u + minute) * 60u + second;
}

static uint8_t FatDayOfWeek(uint32_t year, uint32_t month, uint32_t day) {
    uint64_t epoch = FatCalendarToEpochSeconds(year, month, day, 0, 0, 0);
    return (uint8_t)(((epoch / 86400u) + 4u) % 7u);
}

static uint16_t FatEncodeDate(const SYS_SYSTEM_TIME* t) {
    uint32_t year = t && t->Year ? t->Year : 1980u;
    uint32_t month = t && t->Month ? t->Month : 1u;
    uint32_t day = t && t->Day ? t->Day : 1u;
    if (year < 1980u) year = 1980u;
    if (year > 2107u) year = 2107u;
    if (month < 1u) month = 1u;
    if (month > 12u) month = 12u;
    uint32_t dim = FatDaysInMonth(year, month);
    if (day < 1u) day = 1u;
    if (day > dim) day = dim;
    return (uint16_t)(((year - 1980u) << 9) | (month << 5) | day);
}

static uint16_t FatEncodeTime(const SYS_SYSTEM_TIME* t) {
    uint32_t hour = t ? t->Hour : 0u;
    uint32_t minute = t ? t->Minute : 0u;
    uint32_t second = t ? t->Second : 0u;
    if (hour > 23u) hour = 23u;
    if (minute > 59u) minute = 59u;
    if (second > 59u) second = 59u;
    return (uint16_t)((hour << 11) | (minute << 5) | (second / 2u));
}

static uint8_t FatEncodeCreateTimeTenth(const SYS_SYSTEM_TIME* t) {
    uint32_t second = t ? t->Second : 0u;
    uint32_t ms = t ? t->Millisecond : 0u;
    if (second > 59u) second = 59u;
    if (ms > 999u) ms = 999u;
    uint32_t v = ((second & 1u) * 100u) + (ms / 10u);
    if (v > 199u) v = 199u;
    return (uint8_t)v;
}

void FatQueryEncodedTime(FAT_ENCODED_TIME* out) {
    SYS_SYSTEM_TIME now;
    memset(&now, 0, sizeof(now));
    if (KeQuerySystemTime(&now) != NTLSTATUS_SUCCESS) {
        now.Year = 1980;
        now.Month = 1;
        now.Day = 1;
    }
    out->Date = FatEncodeDate(&now);
    out->Time = FatEncodeTime(&now);
    out->CreateTimeTenth = FatEncodeCreateTimeTenth(&now);
}

void FatDecodeDateTime(uint16_t date, uint16_t time, uint8_t createTenth,
                              SYS_SYSTEM_TIME* out) {
    memset(out, 0, sizeof(*out));
    uint32_t year = 1980u + ((date >> 9) & 0x7Fu);
    uint32_t month = (date >> 5) & 0x0Fu;
    uint32_t day = date & 0x1Fu;
    uint32_t hour = (time >> 11) & 0x1Fu;
    uint32_t minute = (time >> 5) & 0x3Fu;
    uint32_t second = (time & 0x1Fu) * 2u;

    if (date == 0 || month < 1u || month > 12u) { year = 1980u; month = 1u; day = 1u; }
    uint32_t dim = FatDaysInMonth(year, month);
    if (day < 1u) day = 1u;
    if (day > dim) day = dim;
    if (hour > 23u) hour = 23u;
    if (minute > 59u) minute = 59u;
    if (second > 58u) second = 58u;
    if (createTenth >= 100u && second < 59u) second++;

    out->Year = (uint16_t)year;
    out->Month = (uint8_t)month;
    out->Day = (uint8_t)day;
    out->Hour = (uint8_t)hour;
    out->Minute = (uint8_t)minute;
    out->Second = (uint8_t)second;
    out->Millisecond = (uint16_t)((createTenth % 100u) * 10u);
    out->DayOfWeek = FatDayOfWeek(year, month, day);
    out->Flags = SYS_TIME_FLAG_VALID;
    out->EpochSeconds = FatCalendarToEpochSeconds(year, month, day, hour, minute, second);
}

void FatDecodeAccessDate(uint16_t date, SYS_SYSTEM_TIME* out) {
    FatDecodeDateTime(date, 0, 0, out);
}

KSTATUS FatDirentUpdateMetadata(FAT_VOL* v, uint64_t lba,
                                    uint32_t offset,
                                    const FAT_DIRENT_UPDATE* update) {
    if (!v || !update || !update->Flags) return NTLSTATUS_INVALID_PARAMETER;
    uint8_t* sec = (uint8_t*)MmAllocatePool(v->BytesPerSector);
    if (!sec) return NTLSTATUS_NO_MEMORY;
    KSTATUS s = FatReadSectors(v, lba, 1, sec);
    if (KSUCCESS(s)) {
        DIRENT* d = (DIRENT*)(sec + offset);
        if (update->Flags & FAT_DIRENT_UPDATE_SIZE)
            d->FileSize = update->Size;
        if (update->Flags & FAT_DIRENT_UPDATE_FIRST_CLUSTER)
            FatSetDirentFirstCluster(v, d, update->FirstCluster);
        if (update->Flags & FAT_DIRENT_UPDATE_CREATE_TIME) {
            d->CrtTimeTenth = update->Timestamp.CreateTimeTenth;
            d->CrtTime = update->Timestamp.Time;
            d->CrtDate = update->Timestamp.Date;
        }
        if (update->Flags & FAT_DIRENT_UPDATE_WRITE_TIME) {
            d->WrtTime = update->Timestamp.Time;
            d->WrtDate = update->Timestamp.Date;
        }
        if (update->Flags & FAT_DIRENT_UPDATE_ACCESS_DATE)
            d->LastAccDate = update->Timestamp.Date;
        s = FatWriteSectors(v, lba, 1, sec);
    }
    MmFreePool(sec);
    return s;
}

KSTATUS FatDirentUpdateSize(FAT_VOL* v, uint64_t lba,
                                uint32_t offset, uint32_t size) {
    FAT_DIRENT_UPDATE u;
    memset(&u, 0, sizeof(u));
    FatQueryEncodedTime(&u.Timestamp);
    u.Flags = FAT_DIRENT_UPDATE_SIZE | FAT_DIRENT_UPDATE_WRITE_TIME;
    u.Size = size;
    return FatDirentUpdateMetadata(v, lba, offset, &u);
}

KSTATUS FatDirentUpdateFirstCluster(FAT_VOL* v, uint64_t lba,
                                        uint32_t offset, uint32_t cluster) {
    FAT_DIRENT_UPDATE u;
    memset(&u, 0, sizeof(u));
    u.Flags = FAT_DIRENT_UPDATE_FIRST_CLUSTER;
    u.FirstCluster = cluster;
    return FatDirentUpdateMetadata(v, lba, offset, &u);
}