// FAT asynchronous file I/O.

#include "fat_internal.h"

// Native async read/write
// Metadata work (FAT traversal/allocation and partial-cluster read-modify
// preparation) still runs synchronously in the caller/completion context, but
// file data transfers are submitted as disk IRPs and completion is chained from
// the driver's IRP completion callback instead of the generic I/O worker.

typedef struct FAT_ASYNC_RW {
    LIST_ENTRY Link;
    FAT_VOL* Vol;
    FAT_FILE* File;
    uint8_t IsWrite;

    uint64_t Offset;
    uint32_t OffsetInCluster;
    uint32_t Cluster;
    uint32_t LogicalCluster;
    size_t Remaining;
    size_t Done;
    size_t Requested;

    uint8_t* Dst;
    const uint8_t* Src;
    uint8_t* RunBuf;

    uint32_t CurrentNext;
    uint32_t CurrentLast;
    uint32_t CurrentClusters;
    uint32_t CurrentOffset;
    size_t CurrentCopy;

    IRP Irp;
    DEVICE_OBJECT* CurrentDevice;
    IO_FS_ASYNC_COMPLETE Complete;
    void* CompleteContext;
} FAT_ASYNC_RW;

static SPINLOCK gFatAsyncLock = SPINLOCK_INIT;
static LIST_ENTRY gFatAsyncQueue;
static SEMAPHORE* gFatAsyncSemaphore;
static volatile int gFatAsyncInitState;

static void FatAsyncReadSubmitNext(FAT_ASYNC_RW* ctx);
static void FatAsyncWriteSubmitNext(FAT_ASYNC_RW* ctx);

// Implements the async Worker operation.
static void FatAsyncWorker(void* arg) {
    (void)arg;

    while (true) {
        KSTATUS s = KeWaitSemaphore(gFatAsyncSemaphore, KE_WAIT_INFINITE);
        if (!KSUCCESS(s)) {
            KeYield();
            continue;
        }

        uintptr_t flags = KeAcquireSpinlock(&gFatAsyncLock);
        LIST_ENTRY* node = ListPopHead(&gFatAsyncQueue);
        KeReleaseSpinlock(&gFatAsyncLock, flags);
        if (!node) continue;

        FAT_ASYNC_RW* ctx = CONTAINING_RECORD(node, FAT_ASYNC_RW, Link);
        if (ctx->IsWrite)
            FatAsyncWriteSubmitNext(ctx);
        else
            FatAsyncReadSubmitNext(ctx);
    }
}

// Implements the ensure Async Worker Started operation.
static KSTATUS FatEnsureAsyncWorkerStarted(void) {
    uintptr_t flags = KeAcquireSpinlock(&gFatAsyncLock);
    if (gFatAsyncInitState == 2) {
        KeReleaseSpinlock(&gFatAsyncLock, flags);
        return NTLSTATUS_SUCCESS;
    }
    if (gFatAsyncInitState == 1) {
        KeReleaseSpinlock(&gFatAsyncLock, flags);
        while (gFatAsyncInitState == 1)
            KeYield();
        return gFatAsyncInitState == 2 ? NTLSTATUS_SUCCESS : NTLSTATUS_FAILURE;
    }

    gFatAsyncInitState = 1;
    LIST_INIT(&gFatAsyncQueue);
    KeReleaseSpinlock(&gFatAsyncLock, flags);

    KSTATUS s = KeSemaphoreCreate(0, 0x7fffffff, &gFatAsyncSemaphore);
    if (KSUCCESS(s)) {
        THREAD* worker = (THREAD*)0;
        const KE_THREAD_CREATE_PARAMETERS parameters =
            KE_THREAD_CREATE_NAMED_NORMAL(PriorityNormal, "fat.async");
        s = KeCreateKernelThread((uintptr_t)FatAsyncWorker, 0,
                                 &parameters, &worker);
        (void)worker;
    }

    flags = KeAcquireSpinlock(&gFatAsyncLock);
    gFatAsyncInitState = KSUCCESS(s) ? 2 : 0;
    KeReleaseSpinlock(&gFatAsyncLock, flags);
    return s;
}

// Implements the async Queue Continuation operation.
static void FatAsyncQueueContinuation(FAT_ASYNC_RW* ctx) {
    if (!ctx) return;

    uintptr_t flags = KeAcquireSpinlock(&gFatAsyncLock);
    ListInsertTail(&gFatAsyncQueue, &ctx->Link);
    KeReleaseSpinlock(&gFatAsyncLock, flags);
    KeSignalSemaphore(gFatAsyncSemaphore);
}

// Implements the async Finish operation.
static void FatAsyncFinish(FAT_ASYNC_RW* ctx, KSTATUS status, size_t information) {
    if (!ctx) return;
    if (ctx->CurrentDevice) {
        IoDereferenceDeviceObject(ctx->CurrentDevice);
        ctx->CurrentDevice = (DEVICE_OBJECT*)0;
    }
    if (ctx->RunBuf) MmFreePool(ctx->RunBuf);
    IO_FS_ASYNC_COMPLETE complete = ctx->Complete;
    void* completeContext = ctx->CompleteContext;
    MmFreePool(ctx);
    if (complete) complete(completeContext, status, information);
}

// Implements the async Read Irp Complete operation.
static void FatAsyncReadIrpComplete(IRP* irp, void* completeCtx) {
    FAT_ASYNC_RW* ctx = (FAT_ASYNC_RW*)completeCtx;
    if (!ctx) return;
    if (ctx->CurrentDevice) {
        IoDereferenceDeviceObject(ctx->CurrentDevice);
        ctx->CurrentDevice = (DEVICE_OBJECT*)0;
    }

    if (!KSUCCESS(irp->Status)) {
        FatAsyncFinish(ctx, irp->Status, ctx->Done);
        return;
    }

    memcpy(ctx->Dst, ctx->RunBuf + ctx->CurrentOffset, ctx->CurrentCopy);
    ctx->Dst += ctx->CurrentCopy;
    ctx->Done += ctx->CurrentCopy;
    ctx->Remaining -= ctx->CurrentCopy;
    ctx->OffsetInCluster = 0;
    ctx->LogicalCluster += ctx->CurrentClusters;
    ctx->Cluster = ctx->CurrentNext;
    if (ctx->CurrentClusters) {
        uint32_t sectors = ctx->CurrentClusters * ctx->Vol->SectorsPerCluster;
        uint32_t merges = ctx->CurrentClusters - 1u;
        FatStatsAddReadRun(ctx->Vol, sectors, merges);
    }

    FatAsyncQueueContinuation(ctx);
}

// Implements the async Read Submit Next operation.
static void FatAsyncReadSubmitNext(FAT_ASYNC_RW* ctx) {
    if (ctx->Remaining == 0) {
        FatAsyncFinish(ctx, NTLSTATUS_SUCCESS, ctx->Done);
        return;
    }

    if (!FatIsDataCluster(ctx->Vol, ctx->Cluster)) {
        if (ctx->Done < ctx->Requested)
            memset(ctx->Dst, 0, ctx->Requested - ctx->Done);
        FatAsyncFinish(ctx, NTLSTATUS_DEVICE_ERROR, ctx->Done);
        return;
    }

    FAT_VOL* v = ctx->Vol;
    uint32_t runStart = ctx->Cluster;
    uint32_t runSectors = 0;
    uint32_t next = 0;
    if (!FatBuildReadRun(ctx->File, ctx->LogicalCluster, ctx->Cluster,
                           FAT_READ_RUN_MAX_SECTORS, &runSectors, &next)) {
        FatAsyncFinish(ctx, NTLSTATUS_DEVICE_ERROR, ctx->Done);
        return;
    }

    uint32_t runBytes = runSectors * 512u;
    size_t avail = runBytes - ctx->OffsetInCluster;
    size_t copy = avail < ctx->Remaining ? avail : ctx->Remaining;

    ctx->CurrentNext = next;
    ctx->CurrentLast = runStart + (runSectors / v->SectorsPerCluster) - 1u;
    ctx->CurrentClusters = runSectors / v->SectorsPerCluster;
    ctx->CurrentOffset = ctx->OffsetInCluster;
    ctx->CurrentCopy = copy;

    ctx->CurrentDevice = IoReferenceDiskDevice(v->DriveIndex);
    if (!ctx->CurrentDevice) {
        FatAsyncFinish(ctx, NTLSTATUS_NOT_FOUND, ctx->Done);
        return;
    }

    IoBuildReadWriteIrp(&ctx->Irp, IRP_MJ_READ,
                        FatClusterToLba(v, runStart) * 512ull,
                        runSectors * 512u,
                        ctx->RunBuf);
    ctx->Irp.CompletionFn = FatAsyncReadIrpComplete;
    ctx->Irp.CompletionCtx = ctx;

    KSTATUS s = IoSubmitIrp(ctx->CurrentDevice, &ctx->Irp);
    if (s != NTLSTATUS_PENDING) {
        IoDereferenceDeviceObject(ctx->CurrentDevice);
        ctx->CurrentDevice = (DEVICE_OBJECT*)0;
        FatAsyncFinish(ctx, s, ctx->Done);
    }
}

KSTATUS FatReadAsync(IO_VOLUME_PROVIDER* fs, void* fileCtx,
                              uint64_t offset, void* buf, size_t size,
                              IO_FS_ASYNC_COMPLETE complete,
                              void* completeContext) {
    (void)fs;
    if (!complete) return NTLSTATUS_INVALID_PARAMETER;
    if (!fileCtx || (!buf && size)) {
        complete(completeContext, NTLSTATUS_INVALID_PARAMETER, 0);
        return NTLSTATUS_PENDING;
    }

    FAT_FILE* f = (FAT_FILE*)fileCtx;
    FAT_VOL* v = f->Vol;
    if (offset >= f->FileSize || size == 0) {
        complete(completeContext, NTLSTATUS_SUCCESS, 0);
        return NTLSTATUS_PENDING;
    }
    if (offset + size > f->FileSize)
        size = (size_t)(f->FileSize - offset);

    KSTATUS ws = FatEnsureAsyncWorkerStarted();
    if (!KSUCCESS(ws)) {
        complete(completeContext, ws, 0);
        return NTLSTATUS_PENDING;
    }

    FAT_ASYNC_RW* ctx = (FAT_ASYNC_RW*)MmAllocatePool(sizeof(FAT_ASYNC_RW));
    if (!ctx) {
        complete(completeContext, NTLSTATUS_NO_MEMORY, 0);
        return NTLSTATUS_PENDING;
    }
    memset(ctx, 0, sizeof(*ctx));
    LIST_INIT(&ctx->Link);

    ctx->RunBuf = (uint8_t*)MmAllocatePool(FAT_READ_RUN_MAX_SECTORS * 512u);
    if (!ctx->RunBuf) {
        MmFreePool(ctx);
        complete(completeContext, NTLSTATUS_NO_MEMORY, 0);
        return NTLSTATUS_PENDING;
    }

    ctx->Vol = v;
    ctx->File = f;
    ctx->IsWrite = 0;
    ctx->Offset = offset;
    ctx->OffsetInCluster = (uint32_t)(offset % v->BytesPerCluster);
    ctx->LogicalCluster = (uint32_t)(offset / v->BytesPerCluster);
    ctx->Cluster = FatResolveFileCluster(f, ctx->LogicalCluster);
    ctx->Remaining = size;
    ctx->Requested = size;
    ctx->Dst = (uint8_t*)buf;
    ctx->Complete = complete;
    ctx->CompleteContext = completeContext;

    FatAsyncQueueContinuation(ctx);
    return NTLSTATUS_PENDING;
}

// Implements the async Write Irp Complete operation.
static void FatAsyncWriteIrpComplete(IRP* irp, void* completeCtx) {
    FAT_ASYNC_RW* ctx = (FAT_ASYNC_RW*)completeCtx;
    if (!ctx) return;
    if (ctx->CurrentDevice) {
        IoDereferenceDeviceObject(ctx->CurrentDevice);
        ctx->CurrentDevice = (DEVICE_OBJECT*)0;
    }

    if (!KSUCCESS(irp->Status)) {
        FatAsyncFinish(ctx, irp->Status, ctx->Done);
        return;
    }

    ctx->Src += ctx->CurrentCopy;
    ctx->Done += ctx->CurrentCopy;
    ctx->Remaining -= ctx->CurrentCopy;
    ctx->OffsetInCluster = 0;
    ctx->Cluster = ctx->CurrentNext;

    FatAsyncQueueContinuation(ctx);
}

// Implements the async Write Submit Next operation.
static void FatAsyncWriteSubmitNext(FAT_ASYNC_RW* ctx) {
    if (ctx->Remaining == 0) {
        uint64_t newEnd = ctx->Offset + ctx->Done;
        if (ctx->Done > 0) {
            FAT_DIRENT_UPDATE u;
            memset(&u, 0, sizeof(u));
            FatQueryEncodedTime(&u.Timestamp);
            u.Flags = FAT_DIRENT_UPDATE_WRITE_TIME;
            if (newEnd > ctx->File->FileSize) {
                ctx->File->FileSize = (uint32_t)newEnd;
                u.Flags |= FAT_DIRENT_UPDATE_SIZE;
                u.Size = ctx->File->FileSize;
            }
            FatDirentUpdateMetadata(ctx->Vol, ctx->File->DirentLba,
                                 ctx->File->DirentOffset, &u);
        }
        FatAsyncFinish(ctx, NTLSTATUS_SUCCESS, ctx->Done);
        return;
    }

    FAT_VOL* v = ctx->Vol;
    if (!FatIsDataCluster(ctx->Vol, ctx->Cluster)) {
        if (FatIsDataCluster(ctx->Vol, ctx->CurrentLast)) {
            ctx->Cluster = FatAppendCluster(v, ctx->CurrentLast);
            if (!ctx->Cluster) {
                FatAsyncFinish(ctx, NTLSTATUS_NO_MEMORY, ctx->Done);
                return;
            }
        } else {
            FatAsyncFinish(ctx, NTLSTATUS_DEVICE_ERROR, ctx->Done);
            return;
        }
    }

    uint32_t runStart = ctx->Cluster;
    uint32_t runSectors = v->SectorsPerCluster;
    uint32_t last = ctx->Cluster;
    uint32_t next = FatEntry(v, last);

    if (ctx->OffsetInCluster == 0) {
        while (next == last + 1 && FatIsDataCluster(v, next) &&
               runSectors + v->SectorsPerCluster <= FAT_WRITE_RUN_MAX_SECTORS) {
            last = next;
            next = FatEntry(v, last);
            runSectors += v->SectorsPerCluster;
        }
    }

    uint32_t runBytes = runSectors * 512u;
    size_t avail = runBytes - ctx->OffsetInCluster;
    size_t copy = avail < ctx->Remaining ? avail : ctx->Remaining;

    bool partial = (ctx->OffsetInCluster > 0) || (copy < (size_t)runBytes);
    if (partial) {
        KSTATUS rs = FatReadSectors(v, FatClusterToLba(v, runStart), runSectors, ctx->RunBuf);
        if (!KSUCCESS(rs)) {
            FatAsyncFinish(ctx, rs, ctx->Done);
            return;
        }
    }

    memcpy(ctx->RunBuf + ctx->OffsetInCluster, ctx->Src, copy);

    ctx->CurrentNext = next;
    ctx->CurrentLast = last;
    ctx->CurrentOffset = ctx->OffsetInCluster;
    ctx->CurrentCopy = copy;

    ctx->CurrentDevice = IoReferenceDiskDevice(v->DriveIndex);
    if (!ctx->CurrentDevice) {
        FatAsyncFinish(ctx, NTLSTATUS_NOT_FOUND, ctx->Done);
        return;
    }

    IoBuildReadWriteIrp(&ctx->Irp, IRP_MJ_WRITE,
                        FatClusterToLba(v, runStart) * 512ull,
                        runSectors * 512u,
                        ctx->RunBuf);
    ctx->Irp.CompletionFn = FatAsyncWriteIrpComplete;
    ctx->Irp.CompletionCtx = ctx;

    KSTATUS s = IoSubmitIrp(ctx->CurrentDevice, &ctx->Irp);
    if (s != NTLSTATUS_PENDING) {
        IoDereferenceDeviceObject(ctx->CurrentDevice);
        ctx->CurrentDevice = (DEVICE_OBJECT*)0;
        FatAsyncFinish(ctx, s, ctx->Done);
    }
}

KSTATUS FatWriteAsync(IO_VOLUME_PROVIDER* fs, void* fileCtx,
                               uint64_t offset, const void* buf, size_t size,
                               IO_FS_ASYNC_COMPLETE complete,
                               void* completeContext) {
    (void)fs;
    if (!complete) return NTLSTATUS_INVALID_PARAMETER;
    if (!fileCtx || (!buf && size)) {
        complete(completeContext, NTLSTATUS_INVALID_PARAMETER, 0);
        return NTLSTATUS_PENDING;
    }
    if (size == 0) {
        complete(completeContext, NTLSTATUS_SUCCESS, 0);
        return NTLSTATUS_PENDING;
    }

    KSTATUS ws = FatEnsureAsyncWorkerStarted();
    if (!KSUCCESS(ws)) {
        complete(completeContext, ws, 0);
        return NTLSTATUS_PENDING;
    }

    FAT_FILE* f = (FAT_FILE*)fileCtx;
    FAT_VOL* v = f->Vol;
    FatFileInvalidateChainCache(f);

    if (f->StartCluster < 2) {
        uint32_t first = FatAllocCluster(v);
        if (!first) {
            complete(completeContext, NTLSTATUS_NO_MEMORY, 0);
            return NTLSTATUS_PENDING;
        }
        f->StartCluster = first;
        FatDirentUpdateFirstCluster(v, f->DirentLba, f->DirentOffset, first);
    }

    uint32_t clusterIndex = (uint32_t)(offset / v->BytesPerCluster);
    uint32_t cluster = f->StartCluster;
    for (uint32_t i = 0; i < clusterIndex; i++) {
        uint32_t next = FatEntry(v, cluster);
        if (!FatIsDataCluster(v, next)) {
            next = FatAppendCluster(v, cluster);
            if (!next) {
                complete(completeContext, NTLSTATUS_NO_MEMORY, 0);
                return NTLSTATUS_PENDING;
            }
        }
        cluster = next;
    }

    FAT_ASYNC_RW* ctx = (FAT_ASYNC_RW*)MmAllocatePool(sizeof(FAT_ASYNC_RW));
    if (!ctx) {
        complete(completeContext, NTLSTATUS_NO_MEMORY, 0);
        return NTLSTATUS_PENDING;
    }
    memset(ctx, 0, sizeof(*ctx));
    LIST_INIT(&ctx->Link);

    ctx->RunBuf = (uint8_t*)MmAllocatePool(FAT_WRITE_RUN_MAX_SECTORS * 512u);
    if (!ctx->RunBuf) {
        MmFreePool(ctx);
        complete(completeContext, NTLSTATUS_NO_MEMORY, 0);
        return NTLSTATUS_PENDING;
    }

    ctx->Vol = v;
    ctx->File = f;
    ctx->IsWrite = 1;
    ctx->Offset = offset;
    ctx->OffsetInCluster = (uint32_t)(offset % v->BytesPerCluster);
    ctx->Cluster = cluster;
    ctx->Remaining = size;
    ctx->Requested = size;
    ctx->Src = (const uint8_t*)buf;
    ctx->Complete = complete;
    ctx->CompleteContext = completeContext;

    FatAsyncQueueContinuation(ctx);
    return NTLSTATUS_PENDING;
}