#include <ntlite/ddk/module.h>
/*
 * NTLite ext2/ext3/ext4 read/write filesystem driver backed by lwext4.
 *
 * Keep NTLite policy and I/O glue here; third_party/lwext4 stays upstream.
 */

#include <stdint.h>
#include <stddef.h>
#include <stdbool.h>
#include <string.h>
#include <limits.h>

#include <ntlite/ddi/io_driver.h>
#include <ntlite/ddk/ddk.h>
#include <ntlite/ddk/file.h>
#include <ntlite/ddk/memory.h>
#include <ntlite/ddk/debug.h>
#include <ntlite/ddk/list.h>
#include <ntlite/ddk/mutex.h>
#include <ntlite/ddk/time.h>

#include <ext4.h>
#include <ext4_blockdev.h>
#include <ext4_errno.h>
#include <ext4_oflags.h>

#define NTL_EXT4_SECTOR_SIZE 512u
#define NTL_EXT4_ZERO_CHUNK  4096u

#ifndef INT64_MAX
#define INT64_MAX 9223372036854775807LL
#endif

typedef struct NTL_EXT4_VOLUME {
    uint8_t DiskIndex;
    uint64_t PartitionLba;
    uint64_t PartitionSectors;

    struct ext4_blockdev BlockDev;
    struct ext4_blockdev_iface BlockIface;

    char DeviceName[24];
    char MountName[24];

    IO_VOLUME* Volume;
    LIST_ENTRY MountLink;
    bool Mounted;
    bool JournalStarted;
    bool Dead;
} NTL_EXT4_VOLUME;

typedef struct NTL_EXT4_FILE {
    NTL_EXT4_VOLUME* Volume;
    ext4_file File;
    char* Path;
} NTL_EXT4_FILE;

typedef struct NTL_EXT4_DIR {
    NTL_EXT4_VOLUME* Volume;
    ext4_dir Dir;
    char* Path;
} NTL_EXT4_DIR;

static LIST_ENTRY gExt4Mounts;
static uint32_t gExt4MountedCount;
static uint32_t gExt4Serial;

static MUTEX gExt4MountMutex;
static MUTEX gExt4FsMutex;
static SPINLOCK gExt4MutexInitLock = SPINLOCK_INIT;
static int gExt4MutexesReady;

/* lwext4 also defines LIST_INIT for its BSD queue heads.  Avoid that macro
 * here because NTLite LIST_ENTRY uses the Fwd/Bwd intrusive-list layout. */
static void Ext4ListInit(LIST_ENTRY* head) {
    head->Fwd = head;
    head->Bwd = head;
}

static void Ext4InitStaticMutex(MUTEX* mutex) {
    memset(mutex, 0, sizeof(*mutex));
    ObInitializeObjectHeader(&mutex->Header, ObjectTypeMutex);
    mutex->Lock = (SPINLOCK)SPINLOCK_INIT;
    mutex->Owner = (THREAD*)0;
    Ext4ListInit(&mutex->WaitList);
    Ext4ListInit(&mutex->OwnerLink);
    mutex->DonationClass = KeSchedulingNormal;
    mutex->DonationPriority = PriorityLow;
    mutex->HasDonation = 0;
    mutex->Abandoned = 0;
}

static void Ext4EnsureMutexes(void) {
    if (gExt4MutexesReady)
        return;

    uintptr_t flags = KeAcquireSpinlock(&gExt4MutexInitLock);
    if (!gExt4MutexesReady) {
        Ext4InitStaticMutex(&gExt4MountMutex);
        Ext4InitStaticMutex(&gExt4FsMutex);
        gExt4MutexesReady = 1;
    }
    KeReleaseSpinlock(&gExt4MutexInitLock, flags);
}

static KSTATUS Ext4AcquireFs(void) {
    Ext4EnsureMutexes();
    if (!KeCanWait())
        return NTLSTATUS_INVALID_IRQL;
    KSTATUS status = KeAcquireMutex(&gExt4FsMutex, KE_WAIT_INFINITE);
    return KE_MUTEX_ACQUIRE_SUCCEEDED(status) ? NTLSTATUS_SUCCESS : status;
}

static void Ext4ReleaseFs(void) {
    KeReleaseMutex(&gExt4FsMutex);
}

static KSTATUS Ext4MapError(int error) {
    switch (error) {
        case EOK:       return NTLSTATUS_SUCCESS;
        case ENOENT:
        case ENXIO:
        case ENODEV:    return NTLSTATUS_NOT_FOUND;
        case EEXIST:    return NTLSTATUS_ALREADY_EXISTS;
        case ENOMEM:    return NTLSTATUS_NO_MEMORY;
        case ENOSPC:
        case EFBIG:
        case EMLINK:    return NTLSTATUS_LIMIT_REACHED;
        case EACCES:
        case EPERM:
        case EROFS:     return NTLSTATUS_ACCESS_DENIED;
        case ENOTSUP:   return NTLSTATUS_NOT_SUPPORTED;
        case EIO:       return NTLSTATUS_DEVICE_ERROR;
        case EINVAL:
        case EFAULT:
        case E2BIG:
        case EISDIR:
        case ENOTDIR:
        case ENOTEMPTY: return NTLSTATUS_INVALID_PARAMETER;
        default:        return NTLSTATUS_FAILURE;
    }
}

static NTL_EXT4_VOLUME* Ext4Volume(IO_VOLUME_PROVIDER* provider) {
    return provider ? (NTL_EXT4_VOLUME*)provider->Context : (NTL_EXT4_VOLUME*)0;
}

static bool Ext4PathValid(SYS_STRING_VIEW path, bool allowRoot) {
    if (!SysSvIsValid(path))
        return false;
    if (!allowRoot && path.Length == 0)
        return false;
    return true;
}

/* Converts the VFS-relative path into lwext4's private mount namespace. */
static char* Ext4BuildPath(NTL_EXT4_VOLUME* volume, SYS_STRING_VIEW relative) {
    if (!volume || !SysSvIsValid(relative))
        return (char*)0;

    size_t prefix = strlen(volume->MountName);
    if (relative.Length > SIZE_MAX - prefix - 1)
        return (char*)0;

    size_t bytes = prefix + relative.Length + 1;
    char* path = (char*)MmAllocatePool(bytes);
    if (!path)
        return (char*)0;

    memcpy(path, volume->MountName, prefix);
    if (relative.Length)
        memcpy(path + prefix, relative.Data, relative.Length);
    path[prefix + relative.Length] = '\0';
    return path;
}

static uint32_t Ext4NowSeconds(void) {
    uint64_t now = KeQueryUnixTimeSeconds();
    return now > UINT32_MAX ? UINT32_MAX : (uint32_t)now;
}

#define EXT4_TOUCH_ATIME 0x1u
#define EXT4_TOUCH_MTIME 0x2u
#define EXT4_TOUCH_CTIME 0x4u

static void Ext4TouchLocked(const char* path, uint32_t mask) {
    if (!path) return;
    uint32_t now = Ext4NowSeconds();
    if (mask & EXT4_TOUCH_ATIME) (void)ext4_atime_set(path, now);
    if (mask & EXT4_TOUCH_MTIME) (void)ext4_mtime_set(path, now);
    if (mask & EXT4_TOUCH_CTIME) (void)ext4_ctime_set(path, now);
}

static void Ext4SetEpochTime(SYS_SYSTEM_TIME* out, uint32_t seconds) {
    if (!out)
        return;
    memset(out, 0, sizeof(*out));
    out->Flags = SYS_TIME_FLAG_VALID;
    out->TimeZoneMinutes = SYS_TIME_UNSPECIFIED_TZ;
    out->EpochSeconds = seconds;

    uint64_t days = (uint64_t)seconds / 86400ull;
    uint32_t rem = seconds % 86400u;
    out->Hour = (uint8_t)(rem / 3600u);
    out->Minute = (uint8_t)((rem % 3600u) / 60u);
    out->Second = (uint8_t)(rem % 60u);
    out->DayOfWeek = (uint8_t)((days + 4u) % 7u);

    int64_t z = (int64_t)days + 719468ll;
    int64_t era = (z >= 0 ? z : z - 146096ll) / 146097ll;
    uint32_t doe = (uint32_t)(z - era * 146097ll);
    uint32_t yoe = (doe - doe / 1460u + doe / 36524u - doe / 146096u) / 365u;
    int64_t year = (int64_t)yoe + era * 400ll;
    uint32_t doy = doe - (365u * yoe + yoe / 4u - yoe / 100u);
    uint32_t mp = (5u * doy + 2u) / 153u;
    uint32_t day = doy - (153u * mp + 2u) / 5u + 1u;
    uint32_t month = mp + (mp < 10u ? 3u : (uint32_t)-9);
    year += (month <= 2u);
    out->Year = (uint16_t)year;
    out->Month = (uint8_t)month;
    out->Day = (uint8_t)day;
}

static bool Ext4IsDotEntry(const ext4_direntry* entry) {
    if (!entry)
        return false;
    if (entry->name_length == 1 && entry->name[0] == '.')
        return true;
    return entry->name_length == 2 && entry->name[0] == '.' && entry->name[1] == '.';
}

static char* Ext4BuildDirEntryPath(const NTL_EXT4_DIR* dir, const ext4_direntry* entry) {
    if (!dir || !dir->Path || !entry)
        return (char*)0;

    size_t parentLength = strlen(dir->Path);
    size_t nameLength = entry->name_length;
    bool needSlash = parentLength != 0 && dir->Path[parentLength - 1] != '/';
    size_t separator = needSlash ? 1u : 0u;

    if (nameLength > SIZE_MAX - parentLength - separator - 1u)
        return (char*)0;

    size_t bytes = parentLength + separator + nameLength + 1u;
    char* path = (char*)MmAllocatePool(bytes);
    if (!path)
        return (char*)0;

    memcpy(path, dir->Path, parentLength);
    size_t offset = parentLength;
    if (needSlash)
        path[offset++] = '/';
    if (nameLength)
        memcpy(path + offset, entry->name, nameLength);
    path[offset + nameLength] = '\0';
    return path;
}

/* lwext4 v0.0.3 only shrinks in ext4_ftruncate and refuses seeks beyond EOF. */
static int Ext4GrowFileLocked(ext4_file* file, uint64_t target) {
    static const uint8_t zeros[NTL_EXT4_ZERO_CHUNK] = {0};
    uint64_t current = ext4_fsize(file);
    if (target <= current)
        return EOK;
    if (current > (uint64_t)INT64_MAX || target > (uint64_t)INT64_MAX)
        return EFBIG;

    int r = ext4_fseek(file, (int64_t)current, SEEK_SET);
    if (r != EOK)
        return r;

    while (current < target) {
        uint64_t left = target - current;
        size_t chunk = left < sizeof(zeros) ? (size_t)left : sizeof(zeros);
        size_t written = 0;
        r = ext4_fwrite(file, zeros, chunk, &written);
        if (r != EOK)
            return r;
        if (written != chunk)
            return EIO;
        current += written;
    }
    return EOK;
}

/* ---- lwext4 physical block device ---- */

static NTL_EXT4_VOLUME* Ext4BlockVolume(struct ext4_blockdev* bdev) {
    if (!bdev || !bdev->bdif)
        return (NTL_EXT4_VOLUME*)0;
    return (NTL_EXT4_VOLUME*)bdev->bdif->p_user;
}

static int Ext4BlockOpen(struct ext4_blockdev* bdev) {
    return Ext4BlockVolume(bdev) ? EOK : ENODEV;
}

static int Ext4BlockRead(struct ext4_blockdev* bdev, void* buffer,
                         uint64_t blockId, uint32_t blockCount) {
    NTL_EXT4_VOLUME* volume = Ext4BlockVolume(bdev);
    if (!volume || (!buffer && blockCount))
        return EINVAL;
    if (blockId > volume->PartitionSectors ||
        (uint64_t)blockCount > volume->PartitionSectors - blockId)
        return EIO;
    if (blockCount == 0)
        return EOK;

    KSTATUS status = IoReadDiskSectors(volume->DiskIndex,
                                       volume->PartitionLba + blockId,
                                       blockCount,
                                       buffer);
    return KSUCCESS(status) ? EOK : EIO;
}

static int Ext4BlockWrite(struct ext4_blockdev* bdev, const void* buffer,
                          uint64_t blockId, uint32_t blockCount) {
    NTL_EXT4_VOLUME* volume = Ext4BlockVolume(bdev);
    if (!volume || (!buffer && blockCount))
        return EINVAL;
    if (blockId > volume->PartitionSectors ||
        (uint64_t)blockCount > volume->PartitionSectors - blockId)
        return EIO;
    if (blockCount == 0)
        return EOK;

    KSTATUS status = IoWriteDiskSectors(volume->DiskIndex,
                                        volume->PartitionLba + blockId,
                                        blockCount,
                                        buffer);
    return KSUCCESS(status) ? EOK : EIO;
}

static int Ext4BlockClose(struct ext4_blockdev* bdev) {
    NTL_EXT4_VOLUME* volume = Ext4BlockVolume(bdev);
    if (!volume)
        return ENODEV;
    return KSUCCESS(IoFlushDisk(volume->DiskIndex)) ? EOK : EIO;
}

static void Ext4InitBlockDevice(NTL_EXT4_VOLUME* volume) {
    memset(&volume->BlockIface, 0, sizeof(volume->BlockIface));
    memset(&volume->BlockDev, 0, sizeof(volume->BlockDev));

    volume->BlockIface.open = Ext4BlockOpen;
    volume->BlockIface.bread = Ext4BlockRead;
    volume->BlockIface.bwrite = Ext4BlockWrite;
    volume->BlockIface.close = Ext4BlockClose;
    volume->BlockIface.lock = (int (*)(struct ext4_blockdev*))0;
    volume->BlockIface.unlock = (int (*)(struct ext4_blockdev*))0;
    volume->BlockIface.ph_bsize = NTL_EXT4_SECTOR_SIZE;
    volume->BlockIface.ph_bcnt = volume->PartitionSectors;
    volume->BlockIface.p_user = volume;

    volume->BlockDev.bdif = &volume->BlockIface;
    volume->BlockDev.part_offset = 0;
    volume->BlockDev.part_size = volume->PartitionSectors * NTL_EXT4_SECTOR_SIZE;
}

/* ---- IO_VOLUME_PROVIDER operations ---- */

static KSTATUS Ext4Open(IO_VOLUME_PROVIDER* provider, SYS_STRING_VIEW path, void** fileCtx) {
    NTL_EXT4_VOLUME* volume = Ext4Volume(provider);
    if (!volume || !fileCtx || !Ext4PathValid(path, false))
        return NTLSTATUS_INVALID_PARAMETER;

    char* fullPath = Ext4BuildPath(volume, path);
    if (!fullPath)
        return NTLSTATUS_NO_MEMORY;

    NTL_EXT4_FILE* file = (NTL_EXT4_FILE*)MmAllocatePool(sizeof(*file));
    if (!file) {
        MmFreePool(fullPath);
        return NTLSTATUS_NO_MEMORY;
    }
    memset(file, 0, sizeof(*file));
    file->Volume = volume;

    KSTATUS status = Ext4AcquireFs();
    if (KSUCCESS(status)) {
        int r = ext4_fopen2(&file->File, fullPath, O_RDWR);
        status = Ext4MapError(r);
        Ext4ReleaseFs();
    }

    if (!KSUCCESS(status)) {
        MmFreePool(fullPath);
        MmFreePool(file);
        return status;
    }

    file->Path = fullPath;
    *fileCtx = file;
    return NTLSTATUS_SUCCESS;
}

static KSTATUS Ext4Read(IO_VOLUME_PROVIDER* provider, void* fileCtx, uint64_t offset,
                        void* buffer, size_t size, size_t* bytesRead) {
    (void)provider;
    NTL_EXT4_FILE* file = (NTL_EXT4_FILE*)fileCtx;
    if (!file || (!buffer && size) || !bytesRead)
        return NTLSTATUS_INVALID_PARAMETER;
    *bytesRead = 0;
    if (offset > (uint64_t)INT64_MAX)
        return NTLSTATUS_INVALID_PARAMETER;

    KSTATUS status = Ext4AcquireFs();
    if (!KSUCCESS(status))
        return status;

    uint64_t fileSize = ext4_fsize(&file->File);
    if (offset >= fileSize || size == 0) {
        Ext4ReleaseFs();
        return NTLSTATUS_SUCCESS;
    }

    uint64_t available = fileSize - offset;
    size_t toRead = available < (uint64_t)size ? (size_t)available : size;
    int r = ext4_fseek(&file->File, (int64_t)offset, SEEK_SET);
    if (r == EOK)
        r = ext4_fread(&file->File, buffer, toRead, bytesRead);
    if (r == EOK && *bytesRead)
        Ext4TouchLocked(file->Path, EXT4_TOUCH_ATIME);
    status = Ext4MapError(r);
    Ext4ReleaseFs();
    return status;
}

static KSTATUS Ext4Write(IO_VOLUME_PROVIDER* provider, void* fileCtx, uint64_t offset,
                         const void* buffer, size_t size, size_t* bytesWritten) {
    (void)provider;
    NTL_EXT4_FILE* file = (NTL_EXT4_FILE*)fileCtx;
    if (!file || (!buffer && size) || !bytesWritten)
        return NTLSTATUS_INVALID_PARAMETER;
    *bytesWritten = 0;
    if (offset > (uint64_t)INT64_MAX || (uint64_t)size > (uint64_t)INT64_MAX - offset)
        return NTLSTATUS_LIMIT_REACHED;
    if (size == 0)
        return NTLSTATUS_SUCCESS;

    KSTATUS status = Ext4AcquireFs();
    if (!KSUCCESS(status))
        return status;

    int r = EOK;
    uint64_t fileSize = ext4_fsize(&file->File);
    if (offset > fileSize)
        r = Ext4GrowFileLocked(&file->File, offset);
    if (r == EOK)
        r = ext4_fseek(&file->File, (int64_t)offset, SEEK_SET);
    if (r == EOK)
        r = ext4_fwrite(&file->File, buffer, size, bytesWritten);
    if (r == EOK && *bytesWritten)
        Ext4TouchLocked(file->Path, EXT4_TOUCH_MTIME | EXT4_TOUCH_CTIME);

    status = Ext4MapError(r);
    Ext4ReleaseFs();
    return status;
}

static KSTATUS Ext4Create(IO_VOLUME_PROVIDER* provider, SYS_STRING_VIEW path, void** fileCtx) {
    NTL_EXT4_VOLUME* volume = Ext4Volume(provider);
    if (!volume || !fileCtx || !Ext4PathValid(path, false))
        return NTLSTATUS_INVALID_PARAMETER;

    char* fullPath = Ext4BuildPath(volume, path);
    if (!fullPath)
        return NTLSTATUS_NO_MEMORY;

    NTL_EXT4_FILE* file = (NTL_EXT4_FILE*)MmAllocatePool(sizeof(*file));
    if (!file) {
        MmFreePool(fullPath);
        return NTLSTATUS_NO_MEMORY;
    }
    memset(file, 0, sizeof(*file));
    file->Volume = volume;

    KSTATUS status = Ext4AcquireFs();
    if (KSUCCESS(status)) {
        int r = ext4_fopen2(&file->File, fullPath, O_RDWR | O_CREAT | O_EXCL);
        if (r == EOK)
            Ext4TouchLocked(fullPath, EXT4_TOUCH_ATIME | EXT4_TOUCH_MTIME | EXT4_TOUCH_CTIME);
        status = Ext4MapError(r);
        Ext4ReleaseFs();
    }

    if (!KSUCCESS(status)) {
        MmFreePool(fullPath);
        MmFreePool(file);
        return status;
    }
    file->Path = fullPath;
    *fileCtx = file;
    return NTLSTATUS_SUCCESS;
}

static KSTATUS Ext4Delete(IO_VOLUME_PROVIDER* provider, SYS_STRING_VIEW path) {
    NTL_EXT4_VOLUME* volume = Ext4Volume(provider);
    if (!volume || !Ext4PathValid(path, false))
        return NTLSTATUS_INVALID_PARAMETER;
    char* fullPath = Ext4BuildPath(volume, path);
    if (!fullPath)
        return NTLSTATUS_NO_MEMORY;

    KSTATUS status = Ext4AcquireFs();
    if (KSUCCESS(status)) {
        status = Ext4MapError(ext4_fremove(fullPath));
        Ext4ReleaseFs();
    }
    MmFreePool(fullPath);
    return status;
}

static KSTATUS Ext4Rename(IO_VOLUME_PROVIDER* provider, SYS_STRING_VIEW oldPath,
                          SYS_STRING_VIEW newPath) {
    NTL_EXT4_VOLUME* volume = Ext4Volume(provider);
    if (!volume || !Ext4PathValid(oldPath, false) || !Ext4PathValid(newPath, false))
        return NTLSTATUS_INVALID_PARAMETER;

    char* oldFull = Ext4BuildPath(volume, oldPath);
    char* newFull = Ext4BuildPath(volume, newPath);
    if (!oldFull || !newFull) {
        if (oldFull) MmFreePool(oldFull);
        if (newFull) MmFreePool(newFull);
        return NTLSTATUS_NO_MEMORY;
    }

    KSTATUS status = Ext4AcquireFs();
    if (KSUCCESS(status)) {
        int r;
        if (ext4_inode_exist(oldFull, EXT4_DE_DIR) == EOK)
            r = ext4_dir_mv(oldFull, newFull);
        else
            r = ext4_frename(oldFull, newFull);
        status = Ext4MapError(r);
        Ext4ReleaseFs();
    }

    MmFreePool(oldFull);
    MmFreePool(newFull);
    return status;
}

static KSTATUS Ext4CreateDir(IO_VOLUME_PROVIDER* provider, SYS_STRING_VIEW path) {
    NTL_EXT4_VOLUME* volume = Ext4Volume(provider);
    if (!volume || !Ext4PathValid(path, false))
        return NTLSTATUS_INVALID_PARAMETER;
    char* fullPath = Ext4BuildPath(volume, path);
    if (!fullPath)
        return NTLSTATUS_NO_MEMORY;

    KSTATUS status = Ext4AcquireFs();
    if (KSUCCESS(status)) {
        int r = ext4_dir_mk(fullPath);
        if (r == EOK)
            Ext4TouchLocked(fullPath, EXT4_TOUCH_ATIME | EXT4_TOUCH_MTIME | EXT4_TOUCH_CTIME);
        status = Ext4MapError(r);
        Ext4ReleaseFs();
    }
    MmFreePool(fullPath);
    return status;
}

static KSTATUS Ext4RemoveDir(IO_VOLUME_PROVIDER* provider, SYS_STRING_VIEW path) {
    NTL_EXT4_VOLUME* volume = Ext4Volume(provider);
    if (!volume || !Ext4PathValid(path, false))
        return NTLSTATUS_INVALID_PARAMETER;
    char* fullPath = Ext4BuildPath(volume, path);
    if (!fullPath)
        return NTLSTATUS_NO_MEMORY;

    KSTATUS status = Ext4AcquireFs();
    if (!KSUCCESS(status)) {
        MmFreePool(fullPath);
        return status;
    }

    ext4_dir dir;
    memset(&dir, 0, sizeof(dir));
    int r = ext4_dir_open(&dir, fullPath);
    bool empty = true;
    if (r == EOK) {
        const ext4_direntry* entry;
        while ((entry = ext4_dir_entry_next(&dir)) != (const ext4_direntry*)0) {
            if (entry->inode != 0 && !Ext4IsDotEntry(entry)) {
                empty = false;
                break;
            }
        }
        int closeResult = ext4_dir_close(&dir);
        if (r == EOK)
            r = closeResult;
    }
    if (r == EOK && !empty)
        r = ENOTEMPTY;
    if (r == EOK)
        r = ext4_dir_rm(fullPath);

    status = Ext4MapError(r);
    Ext4ReleaseFs();
    MmFreePool(fullPath);
    return status;
}

static KSTATUS Ext4Truncate(IO_VOLUME_PROVIDER* provider, void* fileCtx,
                            uint64_t newSize, uint32_t fileFlags) {
    (void)provider;
    (void)fileFlags;
    NTL_EXT4_FILE* file = (NTL_EXT4_FILE*)fileCtx;
    if (!file)
        return NTLSTATUS_INVALID_PARAMETER;
    if (newSize > (uint64_t)INT64_MAX)
        return NTLSTATUS_LIMIT_REACHED;

    KSTATUS status = Ext4AcquireFs();
    if (!KSUCCESS(status))
        return status;

    uint64_t oldSize = ext4_fsize(&file->File);
    int r = newSize <= oldSize
        ? ext4_ftruncate(&file->File, newSize)
        : Ext4GrowFileLocked(&file->File, newSize);
    if (r == EOK && newSize != oldSize)
        Ext4TouchLocked(file->Path, EXT4_TOUCH_MTIME | EXT4_TOUCH_CTIME);
    status = Ext4MapError(r);
    Ext4ReleaseFs();
    return status;
}

static uint64_t Ext4GetSize(IO_VOLUME_PROVIDER* provider, void* fileCtx) {
    (void)provider;
    NTL_EXT4_FILE* file = (NTL_EXT4_FILE*)fileCtx;
    if (!file)
        return 0;

    if (!KSUCCESS(Ext4AcquireFs()))
        return file->File.fsize;
    uint64_t size = ext4_fsize(&file->File);
    Ext4ReleaseFs();
    return size;
}

static KSTATUS Ext4QueryInfo(IO_VOLUME_PROVIDER* provider, void* fileCtx, SYS_FILE_STAT* out) {
    NTL_EXT4_VOLUME* volume = Ext4Volume(provider);
    NTL_EXT4_FILE* file = (NTL_EXT4_FILE*)fileCtx;
    if (!volume || !file || !out)
        return NTLSTATUS_INVALID_PARAMETER;

    memset(out, 0, sizeof(*out));
    out->Size = sizeof(*out);
    out->Type = SYS_FILE_TYPE_REGULAR;
    out->Attributes = SYS_FILE_ATTR_ARCHIVE;
    out->LinkCount = 1;
    out->FileSize = Ext4GetSize(provider, fileCtx);
    out->FileId = file->File.inode;
    if (volume->Volume)
        out->VolumeId = IoGetVolumeInstanceId(volume->Volume);

    KSTATUS status = Ext4AcquireFs();
    if (KSUCCESS(status)) {
        uint32_t atime = 0, mtime = 0, ctime = 0;
        if (file->Path && ext4_atime_get(file->Path, &atime) == EOK)
            Ext4SetEpochTime(&out->LastAccessTime, atime);
        if (file->Path && ext4_mtime_get(file->Path, &mtime) == EOK)
            Ext4SetEpochTime(&out->LastWriteTime, mtime);
        if (file->Path && ext4_ctime_get(file->Path, &ctime) == EOK) {
            Ext4SetEpochTime(&out->ChangeTime, ctime);
            Ext4SetEpochTime(&out->CreationTime, ctime);
        }
        Ext4ReleaseFs();
    }
    return NTLSTATUS_SUCCESS;
}

static KSTATUS Ext4Stat(IO_VOLUME_PROVIDER* provider, SYS_STRING_VIEW path, SYS_FILE_STAT* out) {
    NTL_EXT4_VOLUME* volume = Ext4Volume(provider);
    if (!volume || !out || !Ext4PathValid(path, true))
        return NTLSTATUS_INVALID_PARAMETER;
    char* fullPath = Ext4BuildPath(volume, path);
    if (!fullPath)
        return NTLSTATUS_NO_MEMORY;

    KSTATUS status = Ext4AcquireFs();
    if (!KSUCCESS(status)) {
        MmFreePool(fullPath);
        return status;
    }

    memset(out, 0, sizeof(*out));
    out->Size = sizeof(*out);
    out->LinkCount = 1;
    if (volume->Volume)
        out->VolumeId = IoGetVolumeInstanceId(volume->Volume);

    bool isDir = ext4_inode_exist(fullPath, EXT4_DE_DIR) == EOK;
    int r = EOK;
    if (isDir) {
        ext4_dir dir;
        memset(&dir, 0, sizeof(dir));
        r = ext4_dir_open(&dir, fullPath);
        if (r == EOK) {
            out->Type = SYS_FILE_TYPE_DIRECTORY;
            out->Attributes = SYS_FILE_ATTR_DIRECTORY;
            out->FileId = dir.f.inode;
            r = ext4_dir_close(&dir);
        }
    } else {
        ext4_file file;
        memset(&file, 0, sizeof(file));
        r = ext4_fopen2(&file, fullPath, O_RDONLY);
        if (r == EOK) {
            out->Type = SYS_FILE_TYPE_REGULAR;
            out->Attributes = SYS_FILE_ATTR_ARCHIVE;
            out->FileId = file.inode;
            out->FileSize = ext4_fsize(&file);
            r = ext4_fclose(&file);
        }
    }

    if (r == EOK) {
        uint32_t atime = 0, mtime = 0, ctime = 0;
        if (ext4_atime_get(fullPath, &atime) == EOK)
            Ext4SetEpochTime(&out->LastAccessTime, atime);
        if (ext4_mtime_get(fullPath, &mtime) == EOK)
            Ext4SetEpochTime(&out->LastWriteTime, mtime);
        if (ext4_ctime_get(fullPath, &ctime) == EOK) {
            Ext4SetEpochTime(&out->ChangeTime, ctime);
            Ext4SetEpochTime(&out->CreationTime, ctime);
        }
    }

    status = Ext4MapError(r);
    Ext4ReleaseFs();
    MmFreePool(fullPath);
    return status;
}

static KSTATUS Ext4QueryIdentity(IO_VOLUME_PROVIDER* provider, void* fileCtx,
                                 IO_FILE_IDENTITY* out) {
    NTL_EXT4_VOLUME* volume = Ext4Volume(provider);
    NTL_EXT4_FILE* file = (NTL_EXT4_FILE*)fileCtx;
    if (!volume || !file || !out)
        return NTLSTATUS_INVALID_PARAMETER;

    memset(out, 0, sizeof(*out));
    if (volume->Volume)
        out->VolumeId = IoGetVolumeInstanceId(volume->Volume);
    out->FileId = file->File.inode;
    out->Size = Ext4GetSize(provider, fileCtx);
    return NTLSTATUS_SUCCESS;
}

static KSTATUS Ext4Flush(IO_VOLUME_PROVIDER* provider, void* fileCtx) {
    (void)fileCtx;
    NTL_EXT4_VOLUME* volume = Ext4Volume(provider);
    if (!volume)
        return NTLSTATUS_INVALID_PARAMETER;

    KSTATUS status = Ext4AcquireFs();
    if (!KSUCCESS(status))
        return status;
    int r = ext4_cache_flush(volume->MountName);
    Ext4ReleaseFs();
    if (r != EOK)
        return Ext4MapError(r);
    return IoFlushDisk(volume->DiskIndex);
}

static void Ext4Close(IO_VOLUME_PROVIDER* provider, void* fileCtx) {
    (void)provider;
    NTL_EXT4_FILE* file = (NTL_EXT4_FILE*)fileCtx;
    if (!file)
        return;
    if (KSUCCESS(Ext4AcquireFs())) {
        (void)ext4_fclose(&file->File);
        Ext4ReleaseFs();
    }
    if (file->Path) MmFreePool(file->Path);
    MmFreePool(file);
}

static void Ext4QueryStats(IO_VOLUME_PROVIDER* provider, IO_FILESYSTEM_STATS* out) {
    (void)provider;
    if (out)
        memset(out, 0, sizeof(*out));
}

static KSTATUS Ext4OpenDir(IO_VOLUME_PROVIDER* provider, SYS_STRING_VIEW path, void** dirCtx) {
    NTL_EXT4_VOLUME* volume = Ext4Volume(provider);
    if (!volume || !dirCtx || !Ext4PathValid(path, true))
        return NTLSTATUS_INVALID_PARAMETER;
    char* fullPath = Ext4BuildPath(volume, path);
    if (!fullPath)
        return NTLSTATUS_NO_MEMORY;

    NTL_EXT4_DIR* dir = (NTL_EXT4_DIR*)MmAllocatePool(sizeof(*dir));
    if (!dir) {
        MmFreePool(fullPath);
        return NTLSTATUS_NO_MEMORY;
    }
    memset(dir, 0, sizeof(*dir));
    dir->Volume = volume;

    KSTATUS status = Ext4AcquireFs();
    if (KSUCCESS(status)) {
        status = Ext4MapError(ext4_dir_open(&dir->Dir, fullPath));
        Ext4ReleaseFs();
    }

    if (!KSUCCESS(status)) {
        MmFreePool(fullPath);
        MmFreePool(dir);
        return status;
    }
    dir->Path = fullPath;
    *dirCtx = dir;
    return NTLSTATUS_SUCCESS;
}

static KSTATUS Ext4ReadDir(IO_VOLUME_PROVIDER* provider, void* dirCtx, KDIR_ENTRY* out) {
    (void)provider;
    NTL_EXT4_DIR* dir = (NTL_EXT4_DIR*)dirCtx;
    if (!dir || !out)
        return NTLSTATUS_INVALID_PARAMETER;

    KSTATUS status = Ext4AcquireFs();
    if (!KSUCCESS(status))
        return status;

    const ext4_direntry* entry;
    do {
        entry = ext4_dir_entry_next(&dir->Dir);
    } while (entry && (entry->inode == 0 || Ext4IsDotEntry(entry)));

    if (!entry) {
        Ext4ReleaseFs();
        return NTLSTATUS_NO_MORE_FILES;
    }

    memset(out, 0, sizeof(*out));
    // ext4 allows arbitrary non-NUL filename bytes. Project malformed UTF-8
    // to U+FFFD only for enumeration; keep raw bytes for ext4 metadata lookups.
    // The displayed name is not a reversible path for malformed disk names.
    (void)SysSvCopyUtf8WithReplacement(
        SysSvFromParts(entry->name, entry->name_length), out->Name, sizeof(out->Name));
    out->Attr = entry->inode_type == EXT4_DE_DIR
        ? SYS_FILE_ATTR_DIRECTORY
        : SYS_FILE_ATTR_ARCHIVE;

    /* SYS_DIR_ENTRY carries the size directly.  NTFS/FAT populate it during
     * enumeration, so do the same here instead of making `dir` report 0. */
    char* childPath = Ext4BuildDirEntryPath(dir, entry);
    if (childPath) {
        if (entry->inode_type != EXT4_DE_DIR) {
            ext4_file file;
            memset(&file, 0, sizeof(file));
            if (ext4_fopen2(&file, childPath, O_RDONLY) == EOK) {
                uint64_t fileSize = ext4_fsize(&file);
                out->FileSize = fileSize > UINT32_MAX ? UINT32_MAX : (uint32_t)fileSize;
                (void)ext4_fclose(&file);
            }
        }

        uint32_t timestamp = 0;
        if (ext4_ctime_get(childPath, &timestamp) == EOK)
            Ext4SetEpochTime(&out->CreateTime, timestamp);
        if (ext4_atime_get(childPath, &timestamp) == EOK)
            Ext4SetEpochTime(&out->LastAccessTime, timestamp);
        if (ext4_mtime_get(childPath, &timestamp) == EOK)
            Ext4SetEpochTime(&out->LastWriteTime, timestamp);
        MmFreePool(childPath);
    }

    Ext4ReleaseFs();
    return NTLSTATUS_SUCCESS;
}

static void Ext4CloseDir(IO_VOLUME_PROVIDER* provider, void* dirCtx) {
    (void)provider;
    NTL_EXT4_DIR* dir = (NTL_EXT4_DIR*)dirCtx;
    if (!dir)
        return;
    if (KSUCCESS(Ext4AcquireFs())) {
        (void)ext4_dir_close(&dir->Dir);
        Ext4ReleaseFs();
    }
    if (dir->Path)
        MmFreePool(dir->Path);
    MmFreePool(dir);
}

static void Ext4VolumeDead(IO_VOLUME_PROVIDER* provider, IO_VOLUME_REMOVAL_REASON reason) {
    (void)reason;
    NTL_EXT4_VOLUME* volume = Ext4Volume(provider);
    if (volume)
        volume->Dead = true;
}

static void Ext4DestroyProvider(IO_VOLUME_PROVIDER* provider) {
    NTL_EXT4_VOLUME* volume = Ext4Volume(provider);
    if (!volume)
        return;

    bool canFree = true;
    if (KSUCCESS(Ext4AcquireFs())) {
        if (volume->Mounted) {
            if (volume->JournalStarted)
                (void)ext4_journal_stop(volume->MountName);
            (void)ext4_cache_flush(volume->MountName);
            int r = ext4_umount(volume->MountName);
            if (r == EOK) {
                volume->Mounted = false;
                /* SharpExt4 v0.0.3 unregister clears the slot but returns ENOENT. */
                (void)ext4_device_unregister(volume->DeviceName);
            } else {
                /* Static lwext4 tables may still reference this blockdev. Leak instead of UAF. */
                canFree = false;
                DbgPrint("EXT4: unmount failed for %s rc=%d; retaining mount context\n",
                         volume->MountName, r);
            }
        }
        Ext4ReleaseFs();
    } else {
        canFree = false;
    }

    if (canFree) {
        memset(volume, 0, sizeof(*volume));
        MmFreePool(volume);
    }
}

static void Ext4InitProvider(IO_VOLUME_PROVIDER* provider, NTL_EXT4_VOLUME* volume) {
    memset(provider, 0, sizeof(*provider));
    provider->Open = Ext4Open;
    provider->Read = Ext4Read;
    provider->Write = Ext4Write;
    provider->Create = Ext4Create;
    provider->Delete = Ext4Delete;
    provider->Rename = Ext4Rename;
    provider->CreateDir = Ext4CreateDir;
    provider->RemoveDir = Ext4RemoveDir;
    provider->Truncate = Ext4Truncate;
    provider->GetSize = Ext4GetSize;
    provider->QueryInfo = Ext4QueryInfo;
    provider->Stat = Ext4Stat;
    provider->QueryIdentity = Ext4QueryIdentity;
    provider->Flush = Ext4Flush;
    provider->Close = Ext4Close;
    provider->QueryStats = Ext4QueryStats;
    provider->OpenDir = Ext4OpenDir;
    provider->ReadDir = Ext4ReadDir;
    provider->CloseDir = Ext4CloseDir;
    provider->VolumeDead = Ext4VolumeDead;
    provider->Destroy = Ext4DestroyProvider;
    provider->Context = volume;
}

/* ---- partition discovery / mount lifecycle ---- */

static NTL_EXT4_VOLUME* Ext4FindMountLocked(const IO_PARTITION_INFO* partition) {
    if (!partition)
        return (NTL_EXT4_VOLUME*)0;
    for (LIST_ENTRY* entry = gExt4Mounts.Fwd; entry != &gExt4Mounts; entry = entry->Fwd) {
        NTL_EXT4_VOLUME* volume = CONTAINING_RECORD(entry, NTL_EXT4_VOLUME, MountLink);
        if (volume->DiskIndex == partition->DiskIndex &&
            volume->PartitionLba == partition->StartLba)
            return volume;
    }
    return (NTL_EXT4_VOLUME*)0;
}

static bool Ext4ProbePartition(const IO_PARTITION_INFO* partition) {
    if (!partition || partition->SectorCount < 3)
        return false;
    uint8_t sector[NTL_EXT4_SECTOR_SIZE];
    KSTATUS status = IoReadDiskSectors(partition->DiskIndex,
                                       partition->StartLba + 2,
                                       1,
                                       sector);
    if (!KSUCCESS(status))
        return false;
    /* ext superblock starts at byte 1024; s_magic is +56. */
    return sector[56] == 0x53 && sector[57] == 0xEF;
}

static void Ext4Hex8(char* out, uint32_t value) {
    static const char digits[] = "0123456789abcdef";
    for (int i = 7; i >= 0; --i) {
        out[i] = digits[value & 0xFu];
        value >>= 4;
    }
}

static void Ext4AssignNames(NTL_EXT4_VOLUME* volume) {
    uint32_t serial = ++gExt4Serial;
    memcpy(volume->DeviceName, "ntlext4", 7);
    Ext4Hex8(volume->DeviceName + 7, serial);
    volume->DeviceName[15] = '\0';

    volume->MountName[0] = '/';
    volume->MountName[1] = 'e';
    Ext4Hex8(volume->MountName + 2, serial);
    volume->MountName[10] = '/';
    volume->MountName[11] = '\0';
}

static bool Ext4TryMountPartition(const IO_PARTITION_INFO* partition) {
    if (!partition || partition->SectorCount < 3 ||
        partition->SectorCount > UINT64_MAX / NTL_EXT4_SECTOR_SIZE ||
        !IoDiskDevicePresent(partition->DiskIndex))
        return false;
    if (!Ext4ProbePartition(partition))
        return false;

    Ext4EnsureMutexes();
    if (!KE_MUTEX_ACQUIRE_SUCCEEDED(KeAcquireMutex(&gExt4MountMutex, KE_WAIT_INFINITE)))
        return false;

    if (Ext4FindMountLocked(partition)) {
        KeReleaseMutex(&gExt4MountMutex);
        return false;
    }

    NTL_EXT4_VOLUME* volume = (NTL_EXT4_VOLUME*)MmAllocatePool(sizeof(*volume));
    if (!volume) {
        KeReleaseMutex(&gExt4MountMutex);
        return false;
    }
    memset(volume, 0, sizeof(*volume));
    Ext4ListInit(&volume->MountLink);
    volume->DiskIndex = partition->DiskIndex;
    volume->PartitionLba = partition->StartLba;
    volume->PartitionSectors = partition->SectorCount;
    Ext4AssignNames(volume);
    Ext4InitBlockDevice(volume);

    KSTATUS status = Ext4AcquireFs();
    int r = EIO;
    if (KSUCCESS(status)) {
        r = ext4_device_register(&volume->BlockDev, volume->DeviceName);
        if (r == EOK) {
            r = ext4_mount(volume->DeviceName, volume->MountName, false);
            if (r == EOK) {
                volume->Mounted = true;
                int recover = ext4_recover(volume->MountName);
                if (recover != EOK && recover != ENOTSUP)
                    DbgPrint("EXT4: journal recovery %s rc=%d (continuing)\n",
                             volume->MountName, recover);
                int journal = ext4_journal_start(volume->MountName);
                if (journal == EOK)
                    volume->JournalStarted = true;
                else
                    DbgPrint("EXT4: journal start %s rc=%d (continuing)\n",
                             volume->MountName, journal);
            } else {
                (void)ext4_device_unregister(volume->DeviceName);
            }
        }
        Ext4ReleaseFs();
    }

    if (r != EOK) {
        DbgPrint("EXT4: disk=%u lba=%llu lwext4 mount failed rc=%d\n",
                 partition->DiskIndex, (unsigned long long)partition->StartLba, r);
        MmFreePool(volume);
        KeReleaseMutex(&gExt4MountMutex);
        return false;
    }

    IO_VOLUME_PROVIDER provider;
    Ext4InitProvider(&provider, volume);
    IO_VOLUME_SOURCE source;
    source.Type = IoVolumeSourcePartition;
    source.DiskIndex = volume->DiskIndex;
    memset(source.Reserved, 0, sizeof(source.Reserved));
    source.StartLba = volume->PartitionLba;
    source.SectorCount = volume->PartitionSectors;

    IO_VOLUME* ioVolume = (IO_VOLUME*)0;
    status = IoCreateVolume(&provider, &source, &ioVolume);
    if (!KSUCCESS(status)) {
        if (ioVolume) {
            (void)IoRemoveVolume(ioVolume, IoVolumeUnmounted);
            IoReleaseVolume(ioVolume);
        } else {
            /* No IO_VOLUME owns the provider yet, so clean it here. */
            Ext4DestroyProvider(&provider);
        }
        DbgPrint("EXT4: disk=%u lba=%llu IoCreateVolume failed status=0x%x\n",
                 partition->DiskIndex, (unsigned long long)partition->StartLba, status);
        KeReleaseMutex(&gExt4MountMutex);
        return false;
    }

    volume->Volume = ioVolume;
    ListInsertTail(&gExt4Mounts, &volume->MountLink);
    gExt4MountedCount++;
    uint32_t instance = IoGetVolumeInstanceId(ioVolume);
    uint32_t count = gExt4MountedCount;

    DbgPrint("EXT4: disk=%u lba=%llu mounted volume=%u at %s; total=%u\n",
             volume->DiskIndex, (unsigned long long)volume->PartitionLba,
             instance, volume->MountName, count);
    KeReleaseMutex(&gExt4MountMutex);
    return true;
}

static void Ext4UnmountPartition(const IO_PARTITION_INFO* partition, IO_PARTITION_CHANGE change) {
    if (!partition)
        return;

    Ext4EnsureMutexes();
    if (!KE_MUTEX_ACQUIRE_SUCCEEDED(KeAcquireMutex(&gExt4MountMutex, KE_WAIT_INFINITE)))
        return;

    NTL_EXT4_VOLUME* volume = Ext4FindMountLocked(partition);
    if (!volume) {
        KeReleaseMutex(&gExt4MountMutex);
        return;
    }

    IO_VOLUME* ioVolume = volume->Volume;
    uint32_t instance = IoGetVolumeInstanceId(ioVolume);
    ListRemove(&volume->MountLink);
    volume->Volume = (IO_VOLUME*)0;
    if (gExt4MountedCount)
        gExt4MountedCount--;
    uint32_t count = gExt4MountedCount;
    KeReleaseMutex(&gExt4MountMutex);

    (void)IoRemoveVolume(ioVolume, IoVolumeBackingRemoved);
    IoReleaseVolume(ioVolume);
    DbgPrint("EXT4: disk=%u lba=%llu %s; unmounted volume=%u; total=%u\n",
             partition->DiskIndex, (unsigned long long)partition->StartLba,
             change == IoPartitionSurpriseRemoved ? "surprise removed" : "removed",
             instance, count);
}

static void Ext4PartitionChange(const IO_PARTITION_INFO* partition,
                                IO_PARTITION_CHANGE change,
                                void* context) {
    (void)context;
    switch (change) {
        case IoPartitionArrived:
            (void)Ext4TryMountPartition(partition);
            break;
        case IoPartitionSurpriseRemoved:
        case IoPartitionRemoved:
            Ext4UnmountPartition(partition, change);
            break;
        default:
            break;
    }
}

static KSTATUS DriverInitialize(NTL_MODULE* module) {
    (void)module;
    DbgPrint("EXT4: Init\n");

    Ext4ListInit(&gExt4Mounts);
    gExt4MountedCount = 0;
    gExt4Serial = 0;
    Ext4EnsureMutexes();

    KSTATUS status = IoRegisterPartitionChangeCallback(Ext4PartitionChange, (void*)0);
    if (!KSUCCESS(status)) {
        DbgPrint("EXT4: partition callback registration failed status=0x%x\n", status);
        return status;
    }

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

NTL_DECLARE_MODULE(DriverInitialize);