#include <ntlite/ddk/module.h>
// FAT mount discovery, BPB validation, and driver lifetime.

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

// FAT BPB validation.
typedef struct FAT_BPB_LAYOUT {
    FAT_VARIANT Variant;
    uint32_t TotalSectors;
    uint32_t FatSize;
    uint32_t RootDirSectors;
    uint32_t DataSectors;
    uint32_t TotalClusters;
} FAT_BPB_LAYOUT;

static bool FatComputeBpbLayout(const BPB* bpb, FAT_BPB_LAYOUT* out) {
    if (!bpb || !out) return false;
    memset(out, 0, sizeof(*out));

    if (bpb->BytesPerSector != 512) return false;
    if (bpb->SectorsPerCluster == 0) return false;
    if (bpb->NumFats == 0) return false;
    if (bpb->ReservedSectors == 0) return false;

    uint32_t totalSectors = bpb->TotalSectors16 ? bpb->TotalSectors16 : bpb->TotalSectors32;
    uint32_t fatSize = bpb->FatSize16 ? bpb->FatSize16 : bpb->FatSize32;
    if (!totalSectors || !fatSize) return false;

    uint32_t rootDirSectors = ((uint32_t)bpb->RootEntryCount * sizeof(DIRENT) +
                               (uint32_t)bpb->BytesPerSector - 1u) /
                              (uint32_t)bpb->BytesPerSector;
    uint32_t metadataSectors = (uint32_t)bpb->ReservedSectors +
                               (uint32_t)bpb->NumFats * fatSize +
                               rootDirSectors;
    if (totalSectors <= metadataSectors) return false;

    uint32_t dataSectors = totalSectors - metadataSectors;
    uint32_t totalClusters = dataSectors / bpb->SectorsPerCluster;
    FAT_VARIANT variant;
    if (totalClusters < 4085u)
        variant = FatVariant12;
    else if (totalClusters < 65525u)
        variant = FatVariant16;
    else
        variant = FatVariant32;

    if (variant == FatVariant16) {
        if (bpb->FatSize16 == 0 || bpb->RootEntryCount == 0) return false;
    } else if (variant == FatVariant32) {
        if (bpb->FatSize16 != 0 || bpb->RootEntryCount != 0) return false;
        if (bpb->FatSize32 == 0 || bpb->RootCluster < 2) return false;
    }

    out->Variant = variant;
    out->TotalSectors = totalSectors;
    out->FatSize = fatSize;
    out->RootDirSectors = rootDirSectors;
    out->DataSectors = dataSectors;
    out->TotalClusters = totalClusters;
    return true;
}


// Driver entry and dynamic partition-backed mount instances.
static LIST_ENTRY gFatMounts;
static uint32_t gMountedCount;
static MUTEX gFatMountMutex;
static SPINLOCK gFatMountMutexInitLock = SPINLOCK_INIT;
static int gFatMountMutexReady;

// Marks only this FAT mount's cache state unusable when its backing volume dies.
void FatVolumeDead(IO_VOLUME_PROVIDER* provider, IO_VOLUME_REMOVAL_REASON reason) {
    (void)reason;
    if (!provider || !provider->Context)
        return;
    FatCacheInvalidateVolume((FAT_VOL*)provider->Context);
}

// Releases one dynamically allocated FAT mount after every volume reference has drained.
void FatDestroyVolumeProvider(IO_VOLUME_PROVIDER* provider) {
    if (!provider || !provider->Context)
        return;
    FAT_VOL* volume = (FAT_VOL*)provider->Context;
    FatCacheInvalidateVolume(volume);
    memset(volume, 0, sizeof(*volume));
    MmFreePool(volume);
}

// Initializes the mount-list mutex without allocating kernel objects.
static void FatMountEnsureMutex(void) {
    if (gFatMountMutexReady)
        return;

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

// Returns one active FAT mount for an exact partition extent while the mount mutex is held.
static FAT_VOL* FatFindMountLocked(const IO_PARTITION_INFO* partition) {
    if (!partition)
        return (FAT_VOL*)0;
    for (LIST_ENTRY* entry = gFatMounts.Fwd; entry != &gFatMounts; entry = entry->Fwd) {
        FAT_VOL* volume = CONTAINING_RECORD(entry, FAT_VOL, MountLink);
        if (volume->DriveIndex == partition->DiskIndex &&
            volume->PartitionLba == partition->StartLba) {
            return volume;
        }
    }
    return (FAT_VOL*)0;
}

// Probes one I/O-core partition extent and mounts it when its VBR describes supported FAT16/FAT32.
static bool FatTryMountPartition(const IO_PARTITION_INFO* partition) {
    if (!partition || partition->SectorCount == 0 || !IoDiskDevicePresent(partition->DiskIndex))
        return false;

    FatMountEnsureMutex();
    if (!KE_MUTEX_ACQUIRE_SUCCEEDED(KeAcquireMutex(&gFatMountMutex, KE_WAIT_INFINITE)))
        return false;

    if (FatFindMountLocked(partition)) {
        KeReleaseMutex(&gFatMountMutex);
        return false;
    }

    FAT_VOL* v = (FAT_VOL*)MmAllocatePool(sizeof(FAT_VOL));
    if (!v) {
        KeReleaseMutex(&gFatMountMutex);
        return false;
    }
    memset(v, 0, sizeof(*v));
    LIST_INIT(&v->MountLink);
    v->DriveIndex = partition->DiskIndex;
    v->PartitionLba = partition->StartLba;
    v->PartitionSize = partition->SectorCount;

    uint8_t sector[512];
    if (!KSUCCESS(FatReadSectors(v, v->PartitionLba, 1, sector))) {
        MmFreePool(v);
        KeReleaseMutex(&gFatMountMutex);
        return false;
    }

    BPB* bpb = (BPB*)sector;
    FAT_BPB_LAYOUT layout;
    if (!FatComputeBpbLayout(bpb, &layout) ||
        (layout.Variant != FatVariant16 && layout.Variant != FatVariant32)) {
        MmFreePool(v);
        KeReleaseMutex(&gFatMountMutex);
        return false;
    }

    v->Variant = layout.Variant;
    v->BytesPerSector = bpb->BytesPerSector;
    v->SectorsPerCluster = bpb->SectorsPerCluster;
    v->BytesPerCluster = bpb->BytesPerSector * bpb->SectorsPerCluster;
    v->ReservedSectors = bpb->ReservedSectors;
    v->NumFats = bpb->NumFats;
    v->FatSize = layout.FatSize;
    v->RootEntryCount = bpb->RootEntryCount;
    v->RootDirSectors = layout.RootDirSectors;
    v->RootDirStart = v->PartitionLba + bpb->ReservedSectors +
                      (uint64_t)bpb->NumFats * layout.FatSize;
    v->RootCluster = v->Variant == FatVariant16 ? FAT_FIXED_ROOT_CLUSTER : bpb->RootCluster;
    v->DataStart = v->RootDirStart + layout.RootDirSectors;
    v->NextFreeHint = 2;
    v->TotalClusters = layout.TotalClusters;
    if (v->PartitionSize > layout.TotalSectors)
        v->PartitionSize = layout.TotalSectors;

    FatDebugPrint("FAT: disk=%u lba=%llu %s layout fat=%u root-sectors=%u data=%u clusters=%u\n",
                  v->DriveIndex, (unsigned long long)v->PartitionLba,
                  FatVariantName(v->Variant), v->FatSize,
                  v->RootDirSectors, layout.DataSectors, v->TotalClusters);

    IO_VOLUME_PROVIDER provider;
    FatInitProvider(&provider, v);
    IO_VOLUME_SOURCE volParams;
    volParams.Type = IoVolumeSourcePartition;
    volParams.DiskIndex = v->DriveIndex;
    memset(volParams.Reserved, 0, sizeof(volParams.Reserved));
    volParams.StartLba = v->PartitionLba;
    volParams.SectorCount = v->PartitionSize;

    IO_VOLUME* volume = (IO_VOLUME*)0;
    KSTATUS status = IoCreateVolume(&provider, &volParams, &volume);
    if (!KSUCCESS(status)) {
        if (volume) {
            (void)IoRemoveVolume(volume, IoVolumeUnmounted);
            IoReleaseVolume(volume);
        } else {
            MmFreePool(v);
        }
        DbgPrint("FAT: disk=%u lba=%llu mount failed status=0x%x\n",
                 partition->DiskIndex, (unsigned long long)partition->StartLba, status);
        KeReleaseMutex(&gFatMountMutex);
        return false;
    }

    // Retain the IoCreateVolume caller reference while this provider remains on the active mount list.
    v->Volume = volume;
    ListInsertTail(&gFatMounts, &v->MountLink);
    gMountedCount++;
    uint32_t instanceId = IoGetVolumeInstanceId(volume);
    uint32_t mountedCount = gMountedCount;

    DbgPrint("FAT: disk=%u lba=%llu %s mounted volume=%u; total=%u\n",
             v->DriveIndex, (unsigned long long)v->PartitionLba,
             FatVariantName(v->Variant), instanceId, mountedCount);
    KeReleaseMutex(&gFatMountMutex);
    return true;
}

// Unmounts one exact FAT provider instance when its backing partition disappears.
static void FatUnmountPartition(const IO_PARTITION_INFO* partition, IO_PARTITION_CHANGE change) {
    if (!partition)
        return;

    FatMountEnsureMutex();
    if (!KE_MUTEX_ACQUIRE_SUCCEEDED(KeAcquireMutex(&gFatMountMutex, KE_WAIT_INFINITE)))
        return;

    FAT_VOL* v = FatFindMountLocked(partition);
    if (!v) {
        KeReleaseMutex(&gFatMountMutex);
        return;
    }

    IO_VOLUME* volume = v->Volume;
    uint32_t instanceId = IoGetVolumeInstanceId(volume);
    ListRemove(&v->MountLink);
    v->Volume = (IO_VOLUME*)0;
    if (gMountedCount)
        gMountedCount--;
    uint32_t mountedCount = gMountedCount;
    KeReleaseMutex(&gFatMountMutex);

    (void)IoRemoveVolume(volume, IoVolumeBackingRemoved);
    IoReleaseVolume(volume); // Drop the active mount-list reference.
    DbgPrint("FAT: disk=%u lba=%llu %s; unmounted volume=%u; total=%u\n",
             partition->DiskIndex, (unsigned long long)partition->StartLba,
             change == IoPartitionSurpriseRemoved ? "surprise removed" : "removed",
             instanceId, mountedCount);
}

// Receives immutable partition extents from the I/O core; FAT never parses partition tables.
static void FatPartitionChange(const IO_PARTITION_INFO* partition,
                               IO_PARTITION_CHANGE change,
                               void* context) {
    (void)context;
    switch (change) {
        case IoPartitionArrived:
            (void)FatTryMountPartition(partition);
            break;
        case IoPartitionSurpriseRemoved:
        case IoPartitionRemoved:
            FatUnmountPartition(partition, change);
            break;
        default:
            break;
    }
}

// Initializes FAT and subscribes to the authoritative kernel partition stream.
static KSTATUS DriverInitialize(NTL_MODULE* module) {
    (void)module;
    DbgPrint("FAT: Init\n");


    FatCacheInvalidate();
    LIST_INIT(&gFatMounts);
    gMountedCount = 0;
    FatMountEnsureMutex();

    KSTATUS status = IoRegisterPartitionChangeCallback(FatPartitionChange, NULL);
    if (!KSUCCESS(status)) {
        DbgPrint("FAT: partition callback registration failed status=0x%x\n", status);
        return status;
    }

    if (!gMountedCount)
        DbgPrint("FAT: no mountable FAT partitions found; driver resident for partition changes.\n");
    DbgPrint("FAT: Init complete; mounted %u volume(s).\n", gMountedCount);
    return NTLSTATUS_SUCCESS;
}

NTL_DECLARE_MODULE(DriverInitialize);