/* NTLite block-device backend for libntfs-3g. */

#include "config.h"

#include <stdint.h>
#include <stddef.h>
#include <string.h>
#include <errno.h>
#include <fcntl.h>
#include <stdio.h>
#include <sys/stat.h>

#include <ntlite/ddi/io_driver.h>
#include <ntlite/ddk/memory.h>

#include <ntfs-3g/device.h>
#include "ntfs3g_device.h"

#define NTFS3G_SECTOR_SIZE 512u

static NTL_NTFS3G_DEVICE_CONTEXT* Ntfs3gContext(struct ntfs_device* dev) {
    return dev ? (NTL_NTFS3G_DEVICE_CONTEXT*)dev->d_private : NULL;
}

static int Ntfs3gDeviceOpen(struct ntfs_device* dev, int flags) {
    NTL_NTFS3G_DEVICE_CONTEXT* ctx = Ntfs3gContext(dev);
    if (!ctx) { errno = EINVAL; return -1; }
    if (NDevOpen(dev)) { errno = EBUSY; return -1; }
    if ((flags & O_ACCMODE) != O_RDWR) NDevSetReadOnly(dev);
    else NDevClearReadOnly(dev);
    NDevSetBlock(dev);
    NDevSetOpen(dev);
    ctx->Position = 0;
    return 0;
}

static int Ntfs3gDeviceClose(struct ntfs_device* dev) {
    if (!dev || !NDevOpen(dev)) { errno = EBADF; return -1; }
    NDevClearOpen(dev);
    return 0;
}

static s64 Ntfs3gDeviceSeek(struct ntfs_device* dev, s64 offset, int whence) {
    NTL_NTFS3G_DEVICE_CONTEXT* ctx = Ntfs3gContext(dev);
    if (!ctx) { errno = EINVAL; return -1; }
    int64_t base;
    if (whence == SEEK_SET) base = 0;
    else if (whence == SEEK_CUR) base = ctx->Position;
    else if (whence == SEEK_END) {
        uint64_t bytes = ctx->PartitionSectors * (uint64_t)NTFS3G_SECTOR_SIZE;
        if (bytes > INT64_MAX) { errno = EOVERFLOW; return -1; }
        base = (int64_t)bytes;
    } else { errno = EINVAL; return -1; }
    if ((offset < 0 && base < -offset) || (offset > 0 && base > INT64_MAX - offset)) {
        errno = EOVERFLOW; return -1;
    }
    base += offset;
    if (base < 0) { errno = EINVAL; return -1; }
    ctx->Position = base;
    return base;
}

static int Ntfs3gBounds(NTL_NTFS3G_DEVICE_CONTEXT* ctx, s64 offset, s64 count) {
    if (!ctx || offset < 0 || count < 0) return 0;
    uint64_t bytes = ctx->PartitionSectors * (uint64_t)NTFS3G_SECTOR_SIZE;
    uint64_t uoff = (uint64_t)offset, ucount = (uint64_t)count;
    return uoff <= bytes && ucount <= bytes - uoff;
}

static s64 Ntfs3gPread(struct ntfs_device* dev, void* buffer, s64 count, s64 offset) {
    NTL_NTFS3G_DEVICE_CONTEXT* ctx = Ntfs3gContext(dev);
    if (!buffer && count) { errno = EINVAL; return -1; }
    if (!Ntfs3gBounds(ctx, offset, count)) { errno = EIO; return -1; }
    if (!count) return 0;

    uint8_t* out = (uint8_t*)buffer;
    uint64_t pos = (uint64_t)offset;
    uint64_t left = (uint64_t)count;
    uint8_t* bounce = NULL;

    while (left) {
        uint64_t sector = pos / NTFS3G_SECTOR_SIZE;
        uint32_t inSector = (uint32_t)(pos % NTFS3G_SECTOR_SIZE);
        if (!inSector && left >= NTFS3G_SECTOR_SIZE) {
            uint64_t whole = left / NTFS3G_SECTOR_SIZE;
            uint32_t batch = whole > 0x10000u ? 0x10000u : (uint32_t)whole;
            KSTATUS s = IoReadDiskSectors(ctx->DiskIndex, ctx->PartitionLba + sector, batch, out);
            if (!KSUCCESS(s)) { errno = EIO; if (bounce) MmFreePool(bounce); return -1; }
            uint64_t bytes = (uint64_t)batch * NTFS3G_SECTOR_SIZE;
            pos += bytes; out += bytes; left -= bytes;
            continue;
        }
        if (!bounce) {
            bounce = (uint8_t*)MmAllocatePool(NTFS3G_SECTOR_SIZE);
            if (!bounce) { errno = ENOMEM; return -1; }
        }
        KSTATUS s = IoReadDiskSectors(ctx->DiskIndex, ctx->PartitionLba + sector, 1, bounce);
        if (!KSUCCESS(s)) { errno = EIO; MmFreePool(bounce); return -1; }
        uint32_t chunk = NTFS3G_SECTOR_SIZE - inSector;
        if ((uint64_t)chunk > left) chunk = (uint32_t)left;
        memcpy(out, bounce + inSector, chunk);
        pos += chunk; out += chunk; left -= chunk;
    }
    if (bounce) MmFreePool(bounce);
    return count;
}

static s64 Ntfs3gPwrite(struct ntfs_device* dev, const void* buffer, s64 count, s64 offset) {
    NTL_NTFS3G_DEVICE_CONTEXT* ctx = Ntfs3gContext(dev);
    if (!buffer && count) { errno = EINVAL; return -1; }
    if (NDevReadOnly(dev)) { errno = EROFS; return -1; }
    if (!Ntfs3gBounds(ctx, offset, count)) { errno = EIO; return -1; }
    if (!count) return 0;

    const uint8_t* in = (const uint8_t*)buffer;
    uint64_t pos = (uint64_t)offset;
    uint64_t left = (uint64_t)count;
    uint8_t* bounce = NULL;

    while (left) {
        uint64_t sector = pos / NTFS3G_SECTOR_SIZE;
        uint32_t inSector = (uint32_t)(pos % NTFS3G_SECTOR_SIZE);
        if (!inSector && left >= NTFS3G_SECTOR_SIZE) {
            uint64_t whole = left / NTFS3G_SECTOR_SIZE;
            uint32_t batch = whole > 0x10000u ? 0x10000u : (uint32_t)whole;
            KSTATUS s = IoWriteDiskSectors(ctx->DiskIndex, ctx->PartitionLba + sector, batch, in);
            if (!KSUCCESS(s)) { errno = EIO; if (bounce) MmFreePool(bounce); return -1; }
            uint64_t bytes = (uint64_t)batch * NTFS3G_SECTOR_SIZE;
            pos += bytes; in += bytes; left -= bytes;
            continue;
        }
        if (!bounce) {
            bounce = (uint8_t*)MmAllocatePool(NTFS3G_SECTOR_SIZE);
            if (!bounce) { errno = ENOMEM; return -1; }
        }
        KSTATUS s = IoReadDiskSectors(ctx->DiskIndex, ctx->PartitionLba + sector, 1, bounce);
        if (!KSUCCESS(s)) { errno = EIO; MmFreePool(bounce); return -1; }
        uint32_t chunk = NTFS3G_SECTOR_SIZE - inSector;
        if ((uint64_t)chunk > left) chunk = (uint32_t)left;
        memcpy(bounce + inSector, in, chunk);
        s = IoWriteDiskSectors(ctx->DiskIndex, ctx->PartitionLba + sector, 1, bounce);
        if (!KSUCCESS(s)) { errno = EIO; MmFreePool(bounce); return -1; }
        pos += chunk; in += chunk; left -= chunk;
    }
    if (bounce) MmFreePool(bounce);
    NDevSetDirty(dev);
    return count;
}

static s64 Ntfs3gRead(struct ntfs_device* dev, void* buffer, s64 count) {
    NTL_NTFS3G_DEVICE_CONTEXT* ctx = Ntfs3gContext(dev);
    s64 done = Ntfs3gPread(dev, buffer, count, ctx ? ctx->Position : 0);
    if (done > 0 && ctx) ctx->Position += done;
    return done;
}

static s64 Ntfs3gWrite(struct ntfs_device* dev, const void* buffer, s64 count) {
    NTL_NTFS3G_DEVICE_CONTEXT* ctx = Ntfs3gContext(dev);
    s64 done = Ntfs3gPwrite(dev, buffer, count, ctx ? ctx->Position : 0);
    if (done > 0 && ctx) ctx->Position += done;
    return done;
}

static int Ntfs3gSync(struct ntfs_device* dev) {
    NTL_NTFS3G_DEVICE_CONTEXT* ctx = Ntfs3gContext(dev);
    if (!ctx) { errno = EINVAL; return -1; }
    KSTATUS s = IoFlushDisk(ctx->DiskIndex);
    if (!KSUCCESS(s)) { errno = EIO; return -1; }
    NDevClearDirty(dev);
    return 0;
}

static int Ntfs3gStat(struct ntfs_device* dev, struct stat* st) {
    NTL_NTFS3G_DEVICE_CONTEXT* ctx = Ntfs3gContext(dev);
    if (!ctx || !st) { errno = EINVAL; return -1; }
    memset(st, 0, sizeof(*st));
    st->st_mode = S_IFBLK;
    st->st_size = (off_t)(ctx->PartitionSectors * (uint64_t)NTFS3G_SECTOR_SIZE);
    st->st_blksize = NTFS3G_SECTOR_SIZE;
    st->st_blocks = (blkcnt_t)ctx->PartitionSectors;
    return 0;
}

static int Ntfs3gIoctl(struct ntfs_device* dev, unsigned long request, void* argp) {
    (void)dev; (void)request; (void)argp;
    errno = ENOTSUP;
    return -1;
}

static struct ntfs_device_operations gNtfs3gOps = {
    Ntfs3gDeviceOpen,
    Ntfs3gDeviceClose,
    Ntfs3gDeviceSeek,
    Ntfs3gRead,
    Ntfs3gWrite,
    Ntfs3gPread,
    Ntfs3gPwrite,
    Ntfs3gSync,
    Ntfs3gStat,
    Ntfs3gIoctl,
};

struct ntfs_device* Ntfs3gCreateDevice(NTL_NTFS3G_DEVICE_CONTEXT* context) {
    if (!context) { errno = EINVAL; return NULL; }
    return ntfs_device_alloc("ntlite-ntfs", 0, &gNtfs3gOps, context);
}