// Mounted-volume core. This file owns volume identity, provider lifetime, and dead-volume dispatch semantics.

#include <stdbool.h>
#include <string.h>
#include <kernel/io_volume.h>
#include <kernel/ke_id.h>
#include <kernel/ke_debug.h>
#include <kernel/ob_object.h>
#include <kernel/mm_file_cache.h>
#include <kernel/mm_virt.h>
#include <kernel/ke_event.h>
#include <kernel/ke_sched.h>

static KE_ID_ALLOCATOR gVolumeIds = KE_ID_ALLOCATOR_INIT(1);
static SPINLOCK gVolumeRegistryLock = SPINLOCK_INIT;
static LIST_ENTRY gVolumeRegistry;
static bool gVolumeRegistryInitialized;
static SPINLOCK gVolumeStatsLock = SPINLOCK_INIT;
static IO_VOLUME_STATS gVolumeStats;
static IO_VOLUME_MANAGER_STATS gVolumeManagerStats;
static EVENT* gVolumeChangeEvent;
static SPINLOCK gVolumeChangeLock = SPINLOCK_INIT;
static uint64_t gVolumeGeneration;
static IO_VOLUME_MANAGER gVolumeManagerObject;
static bool gVolumeManagerObjectInitialized;

static void KiIoVolumeAddStat(uint64_t* field, int64_t delta);

// Initializes the dynamic mounted-volume registry on first use.
static void KiIoEnsureVolumeRegistry(void) {
    uintptr_t flags = KeAcquireSpinlock(&gVolumeRegistryLock);
    if (!gVolumeRegistryInitialized) {
        LIST_INIT(&gVolumeRegistry);
        gVolumeRegistryInitialized = 1;
    }
    KeReleaseSpinlock(&gVolumeRegistryLock, flags);
}


// Publishes one mounted-volume registry change to discovery waiters.
static void KiIoNotifyVolumeChange(void) {
    uintptr_t flags = KeAcquireSpinlock(&gVolumeChangeLock);
    gVolumeGeneration++;
    KeReleaseSpinlock(&gVolumeChangeLock, flags);
    if (gVolumeChangeEvent)
        KeSignalEvent(gVolumeChangeEvent);
}

// Initializes the waitable mounted-volume registry notification before providers start.
KSTATUS IoInitializeVolumeManager(void) {
    KiIoEnsureVolumeRegistry();
    if (!gVolumeManagerObjectInitialized) {
        ObInitializeObjectHeader(&gVolumeManagerObject.Header, ObjectTypeVolumeManager);
        gVolumeManagerObjectInitialized = 1;
    }
    if (gVolumeChangeEvent)
        return NTLSTATUS_SUCCESS;
    return KeEventCreate(EventAutoReset, 0, &gVolumeChangeEvent);
}

// References the persistent manager object used to delegate mounted-volume control-plane authority.
KSTATUS IoReferenceVolumeManager(IO_VOLUME_MANAGER** managerOut) {
    if (!managerOut)
        return NTLSTATUS_INVALID_PARAMETER;
    *managerOut = (IO_VOLUME_MANAGER*)0;
    if (!gVolumeManagerObjectInitialized || !ObTryReferenceObject(&gVolumeManagerObject.Header))
        return NTLSTATUS_OBJECT_DEAD;
    *managerOut = &gVolumeManagerObject;
    return NTLSTATUS_SUCCESS;
}

// Captures one stable value snapshot. A generation change during capture causes a retry instead of returning stale metadata as current.
KSTATUS IoSnapshotVolumes(IO_VOLUME_SNAPSHOT_ENTRY** entriesOut, uint32_t* countOut, uint64_t* generationOut) {
    if (!entriesOut || !countOut || !generationOut)
        return NTLSTATUS_INVALID_PARAMETER;
    *entriesOut = (IO_VOLUME_SNAPSHOT_ENTRY*)0;
    *countOut = 0;
    *generationOut = 0;

    KiIoEnsureVolumeRegistry();
    for (;;) {
        uint64_t generationBefore = IoGetVolumeGeneration();
        uintptr_t flags = KeAcquireSpinlock(&gVolumeRegistryLock);
        size_t count = 0;
        LIST_ENTRY* head = &gVolumeRegistry;
        for (LIST_ENTRY* node = head->Fwd; node != head; node = node->Fwd)
            count++;
        KeReleaseSpinlock(&gVolumeRegistryLock, flags);
        if (count > UINT32_MAX)
            return NTLSTATUS_LIMIT_REACHED;

        IO_VOLUME_SNAPSHOT_ENTRY* entries = (IO_VOLUME_SNAPSHOT_ENTRY*)0;
        if (count) {
            entries = (IO_VOLUME_SNAPSHOT_ENTRY*)MmAllocatePool(count * sizeof(*entries));
            if (!entries)
                return NTLSTATUS_NO_MEMORY;
            memset(entries, 0, count * sizeof(*entries));
        }

        size_t used = 0;
        bool overflowed = false;
        flags = KeAcquireSpinlock(&gVolumeRegistryLock);
        for (LIST_ENTRY* node = head->Fwd; node != head; node = node->Fwd) {
            if (used >= count) {
                overflowed = true;
                break;
            }
            IO_VOLUME* volume = CONTAINING_RECORD(node, IO_VOLUME, RegistryLink);
            entries[used].InstanceId = volume->InstanceId;
            used++;
        }
        KeReleaseSpinlock(&gVolumeRegistryLock, flags);

        uint64_t generationAfter = IoGetVolumeGeneration();
        if (!overflowed && generationBefore == generationAfter) {
            *entriesOut = entries;
            *countOut = (uint32_t)used;
            *generationOut = generationAfter;
            KiIoVolumeAddStat(&gVolumeManagerStats.Snapshots, 1);
            return NTLSTATUS_SUCCESS;
        }
        if (entries)
            MmFreePool(entries);
    }
}

// References one currently mounted volume by its runtime instance identity.
KSTATUS IoReferenceVolumeByInstanceId(uint32_t instanceId, IO_VOLUME** volumeOut) {
    if (!instanceId || !volumeOut)
        return NTLSTATUS_INVALID_PARAMETER;
    *volumeOut = (IO_VOLUME*)0;

    KiIoEnsureVolumeRegistry();
    uintptr_t flags = KeAcquireSpinlock(&gVolumeRegistryLock);
    LIST_ENTRY* head = &gVolumeRegistry;
    for (LIST_ENTRY* node = head->Fwd; node != head; node = node->Fwd) {
        IO_VOLUME* volume = CONTAINING_RECORD(node, IO_VOLUME, RegistryLink);
        if (volume->InstanceId != instanceId)
            continue;
        if (!ObTryReferenceObject(&volume->Header)) {
            KeReleaseSpinlock(&gVolumeRegistryLock, flags);
            return NTLSTATUS_OBJECT_DEAD;
        }
        KeReleaseSpinlock(&gVolumeRegistryLock, flags);
        if (!IoVolumeIsMounted(volume)) {
            ObDereferenceObject(&volume->Header);
            return NTLSTATUS_OBJECT_DEAD;
        }
        *volumeOut = volume;
        KiIoVolumeAddStat(&gVolumeManagerStats.VolumeOpens, 1);
        return NTLSTATUS_SUCCESS;
    }
    KeReleaseSpinlock(&gVolumeRegistryLock, flags);
    return NTLSTATUS_NOT_FOUND;
}

// Returns the current dynamic mounted-volume registry generation.
uint64_t IoGetVolumeGeneration(void) {
    uintptr_t flags = KeAcquireSpinlock(&gVolumeChangeLock);
    uint64_t generation = gVolumeGeneration;
    KeReleaseSpinlock(&gVolumeChangeLock, flags);
    return generation;
}

// Waits until the mounted-volume registry changes or the timeout expires.
KSTATUS IoWaitForVolumeChange(uint64_t observedGeneration, uint64_t timeoutMilliseconds) {
    KiIoVolumeAddStat(&gVolumeManagerStats.Waits, 1);
    if (IoGetVolumeGeneration() != observedGeneration) {
        KiIoVolumeAddStat(&gVolumeManagerStats.Wakeups, 1);
        return NTLSTATUS_SUCCESS;
    }
    if (!gVolumeChangeEvent)
        return NTLSTATUS_OBJECT_DEAD;

    EVENT* events[1] = { gVolumeChangeEvent };
    uint32_t selected = 0;
    KSTATUS status = KeWaitForMultipleEvents(events, 1, 0, timeoutMilliseconds, &selected);
    if (KSUCCESS(status))
        KiIoVolumeAddStat(&gVolumeManagerStats.Wakeups, 1);
    return status;
}

// Adds one value to a mounted-volume statistic.
static void KiIoVolumeAddStat(uint64_t* field, int64_t delta) {
    uintptr_t flags = KeAcquireSpinlock(&gVolumeStatsLock);
    if (delta >= 0)
        *field += (uint64_t)delta;
    else {
        uint64_t amount = (uint64_t)(-delta);
        *field = *field > amount ? *field - amount : 0;
    }
    KeReleaseSpinlock(&gVolumeStatsLock, flags);
}

void IoQueryVolumeStats(IO_VOLUME_STATS* out) {
    if (!out) return;
    uintptr_t flags = KeAcquireSpinlock(&gVolumeStatsLock);
    *out = gVolumeStats;
    KeReleaseSpinlock(&gVolumeStatsLock, flags);
}

void IoQueryVolumeManagerStats(IO_VOLUME_MANAGER_STATS* out) {
    if (!out) return;
    uintptr_t flags = KeAcquireSpinlock(&gVolumeStatsLock);
    *out = gVolumeManagerStats;
    KeReleaseSpinlock(&gVolumeStatsLock, flags);
}

// References the unique mounted volume backed by one exact disk extent start.
KSTATUS IoReferenceVolumeByBacking(uint8_t diskIndex, uint64_t startLba, IO_VOLUME** volumeOut) {
    if (!volumeOut)
        return NTLSTATUS_INVALID_PARAMETER;
    *volumeOut = (IO_VOLUME*)0;

    KiIoEnsureVolumeRegistry();
    IO_VOLUME* match = (IO_VOLUME*)0;
    uint32_t matches = 0;
    uintptr_t flags = KeAcquireSpinlock(&gVolumeRegistryLock);
    LIST_ENTRY* head = &gVolumeRegistry;
    for (LIST_ENTRY* node = head->Fwd; node != head; node = node->Fwd) {
        IO_VOLUME* volume = CONTAINING_RECORD(node, IO_VOLUME, RegistryLink);
        if (volume->Source.DiskIndex != diskIndex || volume->Source.StartLba != startLba)
            continue;
        matches++;
        if (!match && ObTryReferenceObject(&volume->Header))
            match = volume;
    }
    KeReleaseSpinlock(&gVolumeRegistryLock, flags);

    if (matches == 0)
        return NTLSTATUS_NOT_FOUND;
    if (matches != 1) {
        if (match)
            ObDereferenceObject(&match->Header);
        return NTLSTATUS_OBJECT_NAME_COLLISION;
    }
    if (!match)
        return NTLSTATUS_OBJECT_DEAD;
    if (!IoVolumeIsMounted(match)) {
        ObDereferenceObject(&match->Header);
        return NTLSTATUS_OBJECT_DEAD;
    }

    *volumeOut = match;
    return NTLSTATUS_SUCCESS;
}

// Returns whether the volume still accepts provider operations.
bool IoVolumeIsMounted(const IO_VOLUME* volume) {
    if (!volume)
        return false;
    IO_VOLUME* mutableVolume = (IO_VOLUME*)volume;
    uintptr_t flags = KeAcquireSpinlock(&mutableVolume->Lock);
    bool mounted = mutableVolume->State == IoVolumeMounted;
    KeReleaseSpinlock(&mutableVolume->Lock, flags);
    return mounted;
}

// Acquires one object reference only while the mounted-volume state remains live.
bool IoVolumeTryReferenceMounted(IO_VOLUME* volume) {
    if (!volume || !ObTryReferenceObject(&volume->Header))
        return false;
    if (IoVolumeIsMounted(volume))
        return true;
    ObDereferenceObject(&volume->Header);
    return false;
}

// Destroys provider state only after namespace, file, directory, and transient references drain.
static void KiIoVolumeDestroy(OBJECT* object) {
    IO_VOLUME* volume = CONTAINING_RECORD(object, IO_VOLUME, Header);
    if (volume->Provider.Destroy)
        volume->Provider.Destroy(&volume->Provider);
    if (volume->BackingDevice) {
        IoDereferenceDeviceObject(volume->BackingDevice);
        volume->BackingDevice = (DEVICE_OBJECT*)0;
    }
    if (volume->InstanceId)
        KiIoVolumeAddStat(&gVolumeStats.ProviderDestroys, 1);
}

// Creates a dynamic mounted volume and inserts it into the mounted-volume registry.
__declspec(dllexport) KSTATUS IoCreateVolume(const IO_VOLUME_PROVIDER* provider,
                                             const IO_VOLUME_SOURCE* source,
                                             IO_VOLUME** volumeOut) {
    if (!provider || !source || !volumeOut || !provider->Open || !provider->Read ||
        !provider->GetSize || !provider->Close || source->SectorCount == 0 ||
        (source->Type != IoVolumeSourcePartition && source->Type != IoVolumeSourceBlockDevice)) {
        return NTLSTATUS_INVALID_PARAMETER;
    }
    if (source->Type == IoVolumeSourceBlockDevice && source->StartLba != 0)
        return NTLSTATUS_INVALID_PARAMETER;
    *volumeOut = (IO_VOLUME*)0;

    DEVICE_OBJECT* backingDevice = IoReferenceDiskDevice(source->DiskIndex);
    if (!backingDevice)
        return NTLSTATUS_NOT_FOUND;
    uint64_t diskSectors = backingDevice->u.Disk.SectorCount;
    if (source->StartLba >= diskSectors || source->SectorCount > diskSectors - source->StartLba ||
        (source->Type == IoVolumeSourceBlockDevice && source->SectorCount != diskSectors)) {
        IoDereferenceDeviceObject(backingDevice);
        return NTLSTATUS_INVALID_PARAMETER;
    }

    IO_VOLUME* volume = (IO_VOLUME*)ObAllocateObject(ObjectTypeVolume, sizeof(IO_VOLUME));
    if (!volume) {
        IoDereferenceDeviceObject(backingDevice);
        return NTLSTATUS_NO_MEMORY;
    }

    uint32_t instanceId = 0;
    KSTATUS status = KeAllocateId32(&gVolumeIds, &instanceId);
    if (!KSUCCESS(status)) {
        IoDereferenceDeviceObject(backingDevice);
        ObDereferenceObject(&volume->Header);
        return status;
    }

    volume->Header.Destroy = KiIoVolumeDestroy;
    volume->Lock = (SPINLOCK)SPINLOCK_INIT;
    volume->InstanceId = instanceId;
    volume->State = IoVolumeMounted;
    volume->RemovalReason = 0;
    volume->Source = *source;
    volume->BackingDevice = backingDevice;
    volume->Provider = *provider;
    LIST_INIT(&volume->RegistryLink);

    KiIoEnsureVolumeRegistry();
    uintptr_t flags = KeAcquireSpinlock(&gVolumeRegistryLock);
    ListInsertTail(&gVolumeRegistry, &volume->RegistryLink);
    volume->RegistryLinked = 1;
    KeReleaseSpinlock(&gVolumeRegistryLock, flags);

    // The registry owns the allocation reference; the creator receives an independent reference.
    ObReferenceObject(&volume->Header);
    *volumeOut = volume;
    KiIoVolumeAddStat(&gVolumeStats.CreatedVolumes, 1);
    KiIoVolumeAddStat(&gVolumeStats.MountedVolumes, 1);
    KiIoNotifyVolumeChange();

    DbgPrint("io: volume %u mounted disk=%u start=%llu sectors=%llu\n",
             instanceId,
             source->DiskIndex,
             (unsigned long long)source->StartLba,
             (unsigned long long)source->SectorCount);
    return NTLSTATUS_SUCCESS;
}

// Marks a mounted volume terminally dead while retaining provider state until references drain.
__declspec(dllexport) KSTATUS IoRemoveVolume(IO_VOLUME* volume, IO_VOLUME_REMOVAL_REASON reason) {
    if (!volume || reason < IoVolumeUnmounted || reason > IoVolumeProviderDead)
        return NTLSTATUS_INVALID_PARAMETER;

    uintptr_t flags = KeAcquireSpinlock(&volume->Lock);
    if (volume->State != IoVolumeMounted) {
        KeReleaseSpinlock(&volume->Lock, flags);
        return NTLSTATUS_OBJECT_DEAD;
    }
    volume->State = IoVolumeDead;
    volume->RemovalReason = (uint32_t)reason;
    KeReleaseSpinlock(&volume->Lock, flags);

    // No future provider operation can start after this point. The callback may cancel
    // provider-specific work, but Context remains valid until KiIoVolumeDestroy.
    if (volume->Provider.VolumeDead)
        volume->Provider.VolumeDead(&volume->Provider, reason);
    MmFileCacheDiscardVolume(volume);

    bool dropRegistryReference = false;
    flags = KeAcquireSpinlock(&gVolumeRegistryLock);
    if (volume->RegistryLinked) {
        ListRemove(&volume->RegistryLink);
        volume->RegistryLinked = 0;
        dropRegistryReference = true;
    }
    KeReleaseSpinlock(&gVolumeRegistryLock, flags);

    KiIoVolumeAddStat(&gVolumeStats.MountedVolumes, -1);
    KiIoVolumeAddStat(&gVolumeStats.RemovedVolumes, 1);
    KiIoNotifyVolumeChange();
    const char* reasonName = reason == IoVolumeUnmounted ? "unmounted" :
                             reason == IoVolumeBackingRemoved ? "backing-removed" : "provider-dead";
    DbgPrint("io: volume %u removed reason=%s\n", volume->InstanceId, reasonName);
    if (dropRegistryReference)
        ObDereferenceObject(&volume->Header);
    return NTLSTATUS_SUCCESS;
}

// Releases a creator-owned mounted-volume reference.
__declspec(dllexport) void IoReleaseVolume(IO_VOLUME* volume) {
    if (volume)
        ObDereferenceObject(&volume->Header);
}

// Returns the monotonic identity of a mounted-volume instance.
__declspec(dllexport) uint32_t IoGetVolumeInstanceId(const IO_VOLUME* volume) {
    return volume ? volume->InstanceId : 0;
}

// Returns the immutable block extent backing this mounted-volume instance.
__declspec(dllexport) KSTATUS IoGetVolumeSource(const IO_VOLUME* volume, IO_VOLUME_SOURCE* sourceOut) {
    if (!volume || !sourceOut)
        return NTLSTATUS_INVALID_PARAMETER;
    *sourceOut = volume->Source;
    return NTLSTATUS_SUCCESS;
}

// Returns whether native provider async read is available while the volume remains mounted.
bool IoVolumeSupportsReadAsync(const IO_VOLUME* volume) {
    return volume && IoVolumeIsMounted(volume) && volume->Provider.ReadAsync != (void*)0;
}

// Returns whether native provider async write is available while the volume remains mounted.
bool IoVolumeSupportsWriteAsync(const IO_VOLUME* volume) {
    return volume && IoVolumeIsMounted(volume) && volume->Provider.WriteAsync != (void*)0;
}

// Validates mounted state before entering provider code.
static KSTATUS KiIoVolumeCheckMounted(IO_VOLUME* volume) {
    if (IoVolumeIsMounted(volume))
        return NTLSTATUS_SUCCESS;
    KiIoVolumeAddStat(&gVolumeStats.DeadIoRejects, 1);
    return NTLSTATUS_OBJECT_DEAD;
}

KSTATUS IoVolumeOpen(IO_VOLUME* volume, SYS_STRING_VIEW path, void** fileCtx) {
    KSTATUS s = KiIoVolumeCheckMounted(volume);
    if (!KSUCCESS(s)) return s;
    return volume->Provider.Open(&volume->Provider, path, fileCtx);
}

KSTATUS IoVolumeRead(IO_VOLUME* volume, void* fileCtx, uint64_t offset,
                     void* buffer, size_t size, size_t* bytesRead) {
    KSTATUS s = KiIoVolumeCheckMounted(volume);
    if (!KSUCCESS(s)) return s;
    if (!volume->Provider.Read) return NTLSTATUS_NOT_SUPPORTED;
    return volume->Provider.Read(&volume->Provider, fileCtx, offset, buffer, size, bytesRead);
}

KSTATUS IoVolumeWrite(IO_VOLUME* volume, void* fileCtx, uint64_t offset,
                      const void* buffer, size_t size, size_t* bytesWritten) {
    KSTATUS s = KiIoVolumeCheckMounted(volume);
    if (!KSUCCESS(s)) return s;
    if (!volume->Provider.Write) return NTLSTATUS_NOT_SUPPORTED;
    return volume->Provider.Write(&volume->Provider, fileCtx, offset, buffer, size, bytesWritten);
}

KSTATUS IoVolumeReadAsync(IO_VOLUME* volume, void* fileCtx, uint64_t offset,
                          void* buffer, size_t size, IO_FS_ASYNC_COMPLETE complete,
                          void* completeContext) {
    KSTATUS s = KiIoVolumeCheckMounted(volume);
    if (!KSUCCESS(s)) return s;
    if (!volume->Provider.ReadAsync) return NTLSTATUS_NOT_SUPPORTED;
    return volume->Provider.ReadAsync(&volume->Provider, fileCtx, offset, buffer, size, complete, completeContext);
}

KSTATUS IoVolumeWriteAsync(IO_VOLUME* volume, void* fileCtx, uint64_t offset,
                           const void* buffer, size_t size, IO_FS_ASYNC_COMPLETE complete,
                           void* completeContext) {
    KSTATUS s = KiIoVolumeCheckMounted(volume);
    if (!KSUCCESS(s)) return s;
    if (!volume->Provider.WriteAsync) return NTLSTATUS_NOT_SUPPORTED;
    return volume->Provider.WriteAsync(&volume->Provider, fileCtx, offset, buffer, size, complete, completeContext);
}

KSTATUS IoVolumeCreate(IO_VOLUME* volume, SYS_STRING_VIEW path, void** fileCtx) {
    KSTATUS s = KiIoVolumeCheckMounted(volume);
    if (!KSUCCESS(s)) return s;
    if (!volume->Provider.Create) return NTLSTATUS_NOT_SUPPORTED;
    return volume->Provider.Create(&volume->Provider, path, fileCtx);
}

KSTATUS IoVolumeDelete(IO_VOLUME* volume, SYS_STRING_VIEW path) {
    KSTATUS s = KiIoVolumeCheckMounted(volume);
    if (!KSUCCESS(s)) return s;
    if (!volume->Provider.Delete) return NTLSTATUS_NOT_SUPPORTED;
    return volume->Provider.Delete(&volume->Provider, path);
}

KSTATUS IoVolumeRename(IO_VOLUME* volume, SYS_STRING_VIEW oldPath, SYS_STRING_VIEW newPath) {
    KSTATUS s = KiIoVolumeCheckMounted(volume);
    if (!KSUCCESS(s)) return s;
    if (!volume->Provider.Rename) return NTLSTATUS_NOT_SUPPORTED;
    return volume->Provider.Rename(&volume->Provider, oldPath, newPath);
}

KSTATUS IoVolumeCreateDir(IO_VOLUME* volume, SYS_STRING_VIEW path) {
    KSTATUS s = KiIoVolumeCheckMounted(volume);
    if (!KSUCCESS(s)) return s;
    if (!volume->Provider.CreateDir) return NTLSTATUS_NOT_SUPPORTED;
    return volume->Provider.CreateDir(&volume->Provider, path);
}

KSTATUS IoVolumeRemoveDir(IO_VOLUME* volume, SYS_STRING_VIEW path) {
    KSTATUS s = KiIoVolumeCheckMounted(volume);
    if (!KSUCCESS(s)) return s;
    if (!volume->Provider.RemoveDir) return NTLSTATUS_NOT_SUPPORTED;
    return volume->Provider.RemoveDir(&volume->Provider, path);
}

KSTATUS IoVolumeTruncate(IO_VOLUME* volume, void* fileCtx, uint64_t newSize, uint32_t fileFlags) {
    KSTATUS s = KiIoVolumeCheckMounted(volume);
    if (!KSUCCESS(s)) return s;
    if (!volume->Provider.Truncate) return NTLSTATUS_NOT_SUPPORTED;
    return volume->Provider.Truncate(&volume->Provider, fileCtx, newSize, fileFlags);
}

KSTATUS IoVolumeGetSize(IO_VOLUME* volume, void* fileCtx, uint64_t* sizeOut) {
    if (!sizeOut) return NTLSTATUS_INVALID_PARAMETER;
    *sizeOut = 0;
    KSTATUS s = KiIoVolumeCheckMounted(volume);
    if (!KSUCCESS(s)) return s;
    if (!volume->Provider.GetSize) return NTLSTATUS_NOT_SUPPORTED;
    *sizeOut = volume->Provider.GetSize(&volume->Provider, fileCtx);
    return NTLSTATUS_SUCCESS;
}

KSTATUS IoVolumeQueryInfo(IO_VOLUME* volume, void* fileCtx, SYS_FILE_STAT* out) {
    KSTATUS s = KiIoVolumeCheckMounted(volume);
    if (!KSUCCESS(s)) return s;
    if (!volume->Provider.QueryInfo) return NTLSTATUS_NOT_SUPPORTED;
    s = volume->Provider.QueryInfo(&volume->Provider, fileCtx, out);
    if (KSUCCESS(s) && out)
        out->VolumeId = volume->InstanceId;
    return s;
}

KSTATUS IoVolumeStat(IO_VOLUME* volume, SYS_STRING_VIEW path, SYS_FILE_STAT* out) {
    KSTATUS s = KiIoVolumeCheckMounted(volume);
    if (!KSUCCESS(s)) return s;
    if (!volume->Provider.Stat) return NTLSTATUS_NOT_SUPPORTED;
    s = volume->Provider.Stat(&volume->Provider, path, out);
    if (KSUCCESS(s) && out)
        out->VolumeId = volume->InstanceId;
    return s;
}

KSTATUS IoVolumeQueryIdentity(IO_VOLUME* volume, void* fileCtx, IO_FILE_IDENTITY* out) {
    KSTATUS s = KiIoVolumeCheckMounted(volume);
    if (!KSUCCESS(s)) return s;
    if (!out || !volume->Provider.QueryIdentity) return NTLSTATUS_NOT_SUPPORTED;
    s = volume->Provider.QueryIdentity(&volume->Provider, fileCtx, out);
    if (KSUCCESS(s))
        out->VolumeId = volume->InstanceId;
    return s;
}

KSTATUS IoVolumeFlush(IO_VOLUME* volume, void* fileCtx) {
    KSTATUS s = KiIoVolumeCheckMounted(volume);
    if (!KSUCCESS(s)) return s;
    if (!volume->Provider.Flush) return NTLSTATUS_SUCCESS;
    return volume->Provider.Flush(&volume->Provider, fileCtx);
}

void IoVolumeClose(IO_VOLUME* volume, void* fileCtx) {
    if (volume && volume->Provider.Close && fileCtx)
        volume->Provider.Close(&volume->Provider, fileCtx);
}

void IoVolumeQueryStats(IO_VOLUME* volume, IO_FILESYSTEM_STATS* out) {
    if (!out) return;
    memset(out, 0, sizeof(*out));
    if (!volume || !IoVolumeIsMounted(volume) || !volume->Provider.QueryStats)
        return;
    volume->Provider.QueryStats(&volume->Provider, out);
}

// Aggregates filesystem traversal statistics across every dynamic mounted volume.
__declspec(dllexport) void IoQueryFilesystemStats(IO_FILESYSTEM_STATS* out) {
    if (!out) return;
    memset(out, 0, sizeof(*out));

    KiIoEnsureVolumeRegistry();
    uintptr_t flags = KeAcquireSpinlock(&gVolumeRegistryLock);
    size_t count = 0;
    LIST_ENTRY* head = &gVolumeRegistry;
    for (LIST_ENTRY* node = head->Fwd; node != head; node = node->Fwd)
        count++;
    KeReleaseSpinlock(&gVolumeRegistryLock, flags);
    if (!count) return;

    IO_VOLUME** snapshot = (IO_VOLUME**)MmAllocatePool(count * sizeof(IO_VOLUME*));
    if (!snapshot) return;
    memset(snapshot, 0, count * sizeof(IO_VOLUME*));

    size_t used = 0;
    flags = KeAcquireSpinlock(&gVolumeRegistryLock);
    for (LIST_ENTRY* node = head->Fwd; node != head && used < count; node = node->Fwd) {
        IO_VOLUME* volume = CONTAINING_RECORD(node, IO_VOLUME, RegistryLink);
        if (ObTryReferenceObject(&volume->Header))
            snapshot[used++] = volume;
    }
    KeReleaseSpinlock(&gVolumeRegistryLock, flags);

    for (size_t i = 0; i < used; ++i) {
        IO_FILESYSTEM_STATS one;
        IoVolumeQueryStats(snapshot[i], &one);
        ObDereferenceObject(&snapshot[i]->Header);

        out->PathLookups += one.PathLookups;
        out->PathCacheHits += one.PathCacheHits;
        out->PathCacheMisses += one.PathCacheMisses;
        out->PathCacheStores += one.PathCacheStores;
        out->PathCacheInvalidations += one.PathCacheInvalidations;
        if (one.PathCacheGeneration > out->PathCacheGeneration)
            out->PathCacheGeneration = one.PathCacheGeneration;
        out->DirectorySearches += one.DirectorySearches;
        out->DirectoryBlocksRead += one.DirectoryBlocksRead;
        out->DirectoryEntriesScanned += one.DirectoryEntriesScanned;
        out->DataReadRuns += one.DataReadRuns;
        out->DataReadSectors += one.DataReadSectors;
        out->DataReadRunMerges += one.DataReadRunMerges;
        if (one.MaxDataReadSpanSectors > out->MaxDataReadSpanSectors)
            out->MaxDataReadSpanSectors = one.MaxDataReadSpanSectors;
    }
    MmFreePool(snapshot);
}

KSTATUS IoVolumeOpenDir(IO_VOLUME* volume, SYS_STRING_VIEW path, void** dirCtx) {
    KSTATUS s = KiIoVolumeCheckMounted(volume);
    if (!KSUCCESS(s)) return s;
    if (!volume->Provider.OpenDir) return NTLSTATUS_NOT_SUPPORTED;
    return volume->Provider.OpenDir(&volume->Provider, path, dirCtx);
}

KSTATUS IoVolumeReadDir(IO_VOLUME* volume, void* dirCtx, KDIR_ENTRY* out) {
    KSTATUS s = KiIoVolumeCheckMounted(volume);
    if (!KSUCCESS(s)) return s;
    if (!volume->Provider.ReadDir) return NTLSTATUS_NOT_SUPPORTED;
    return volume->Provider.ReadDir(&volume->Provider, dirCtx, out);
}

void IoVolumeCloseDir(IO_VOLUME* volume, void* dirCtx) {
    if (volume && volume->Provider.CloseDir && dirCtx)
        volume->Provider.CloseDir(&volume->Provider, dirCtx);
}