// io_file.c - VFS dispatch layer
#include <string.h>
#include <ke_list.h>
#include <ke_process.h>
#include <io_file.h>
#include <mm_virt.h>
#include <ob_object.h>
#include <ke_userenv.h>

#define MAX_DRIVES  8

typedef struct {
    IO_FILESYSTEM   *Fs;
    IO_VOLUME_PARAMS Vol;
} VOLUME_ENTRY;

static VOLUME_ENTRY gVolumes[MAX_DRIVES];

// ---------------------------------------------------------------------------
// Current Working Directory  (drive 0 relative, no trailing backslash)
// e.g. "" = root,  "apps" = root\apps,  "apps\\data" = nested
// ---------------------------------------------------------------------------
#define CWD_MAX 512

__declspec(dllexport) KSTATUS IoGetCwd(char* buf, size_t size) {
    if (!buf || size == 0) return STATUS_INVALID_PARAMETER;

    PROCESS* proc = KeGetCurrentProcess();
    if (!proc) return STATUS_ACCESS_DENIED;

    uintptr_t flags = KeAcquireSpinlock(&proc->CwdLock);
    size_t cwdLen = strlen(proc->CurrentDirectory);
    if (cwdLen + 1 > size) {
        KeReleaseSpinlock(&proc->CwdLock, flags);
        return STATUS_BUFFER_TOO_SMALL;
    }
    memcpy(buf, proc->CurrentDirectory, cwdLen + 1);
    KeReleaseSpinlock(&proc->CwdLock, flags);

    return STATUS_SUCCESS;
}

// Resolve a possibly-relative path component against the CWD.
// Writes the fully-resolved, drive-stripped relative path into 'resolved'
// (caller supplies buf of at least CWD_MAX bytes).
// Returns the drive index.
static uint8_t ResolvePath(const char* path, char* resolved, size_t resolvedSize) {
    // Strip explicit drive prefix "N:\"
    uint8_t drive = 0;
    bool fromDriveRoot = false;
    if (path[0] >= '0' && path[0] <= '7' &&
        path[1] == ':' && (path[2] == '\\' || path[2] == '/')) {
        drive = (uint8_t)(path[0] - '0');
        path  = path + 3;
        fromDriveRoot = true;
    }

    char work[CWD_MAX * 2];
    work[0] = '\0';

    // Initialize our working base
    if (path[0] == '\\' || path[0] == '/' || fromDriveRoot) {
        // Absolute path: start from the root (empty work buffer)
        // The segment loop below will automatically skip the leading separators.
    } else {
        // Relative path: start from the current process's CWD
        PROCESS* proc = KeGetCurrentProcess();
        if (proc) {
            uintptr_t flags = KeAcquireSpinlock(&proc->CwdLock);
            size_t wlen = strlen(proc->CurrentDirectory);
            if (wlen >= sizeof(work)) wlen = sizeof(work) - 1;
            memcpy(work, proc->CurrentDirectory, wlen);
            work[wlen] = '\0';
            KeReleaseSpinlock(&proc->CwdLock, flags);
        }
    }

    // Segment-by-segment canonicalization
    const char* p = path;
    while (*p) {
        // Skip separators
        while (*p == '\\' || *p == '/') p++;
        if (!*p) break;

        // Read next component
        char comp[256];
        int  clen = 0;
        while (*p && *p != '\\' && *p != '/' && clen < 255)
            comp[clen++] = *p++;
        comp[clen] = '\0';

        if (clen == 0) continue;

        if (strcmp(comp, ".") == 0) {
            // Stay in place
        } else if (strcmp(comp, "..") == 0) {
            // Remove last component from work
            int i = (int)strlen(work) - 1;
            while (i >= 0 && work[i] != '\\' && work[i] != '/') i--;
            if (i < 0) i = 0;
            work[i] = '\0';
        } else {
            // Append separator + component
            size_t wl = strlen(work);
            if (wl > 0 && work[wl-1] != '\\' && work[wl-1] != '/') {
                if (wl + 1 < sizeof(work)) { work[wl] = '\\'; work[wl+1] = '\0'; }
            }
            size_t wl2 = strlen(work);
            if (wl2 + clen + 1 < sizeof(work)) {
                memcpy(work + wl2, comp, clen);
                work[wl2 + clen] = '\0';
            }
        }
    }

    // Output the final canonicalized, drive-stripped path
    size_t finalLen = strlen(work);
    if (finalLen + 1 > resolvedSize) finalLen = resolvedSize - 1;
    memcpy(resolved, work, finalLen);
    resolved[finalLen] = '\0';
    
    return drive;
}

__declspec(dllexport) KSTATUS IoSetCwd(const char* path) {
    if (!path) return STATUS_INVALID_PARAMETER;

    PROCESS* proc = KeGetCurrentProcess();
    if (!proc) return STATUS_ACCESS_DENIED;

    // Resolve the new path (this also validates ".." navigation).
    // Note: ResolvePath already uses proc->CurrentDirectory internally!
    char resolved[CWD_MAX];
    uint8_t drive = ResolvePath(path, resolved, sizeof(resolved));

    // Validate: try to open the directory to confirm it actually exists on disk
    IO_FILESYSTEM* fs = gVolumes[drive].Fs;
    if (!fs) return STATUS_NOT_FOUND;
    if (!fs->OpenDir) return STATUS_NOT_SUPPORTED;

    void* dirCtx = NULL;
    KSTATUS s = fs->OpenDir(fs, resolved, &dirCtx);
    if (!KSUCCESS(s)) return s;
    fs->CloseDir(fs, dirCtx);

    // Commit the new CWD to the process block
    size_t len = strlen(resolved);
    if (len + 1 > sizeof(proc->CurrentDirectory)) return STATUS_BUFFER_TOO_SMALL;

    uintptr_t flags = KeAcquireSpinlock(&proc->CwdLock);
    memcpy(proc->CurrentDirectory, resolved, len + 1);

    // Keep the user-visible PEB mirror in sync.  The PEB's CWD buffer is a
    // fixed SYS_CWD_MAX allocation made at process creation, so the pointer
    // remains stable for personalities that cache SYS_PEB*.
    if (proc->UserCwdBase && proc->UserCwdSize) {
        size_t copyLen = len;
        if (copyLen + 1 > proc->UserCwdSize) copyLen = proc->UserCwdSize - 1;
        memcpy((void*)proc->UserCwdBase, resolved, copyLen);
        ((char*)proc->UserCwdBase)[copyLen] = '\0';

        SYS_PEB* peb = (SYS_PEB*)proc->UserPebBase;
        if (peb && peb->Magic == SYS_PEB_MAGIC) {
            peb->CurrentDirectory = (const char*)proc->UserCwdBase;
            peb->CurrentDirectoryLength = copyLen;
        }
    }

    KeReleaseSpinlock(&proc->CwdLock, flags);

    return STATUS_SUCCESS;
}

// ---------------------------------------------------------------------------
// Registration
// ---------------------------------------------------------------------------
__declspec(dllexport) KSTATUS IoRegisterFilesystem(uint8_t volumeIndex,
                              IO_FILESYSTEM *fs,
                              const IO_VOLUME_PARAMS *vol) {
    if (volumeIndex >= MAX_DRIVES || !fs || !vol)
        return STATUS_INVALID_PARAMETER;
    if (gVolumes[volumeIndex].Fs)
        return STATUS_ALREADY_EXISTS;
    gVolumes[volumeIndex].Fs  = fs;
    gVolumes[volumeIndex].Vol = *vol;
    return STATUS_SUCCESS;
}

// ---------------------------------------------------------------------------
// Destructor
// ---------------------------------------------------------------------------
static void FileDestroy(OBJECT* obj) {
    FILE_OBJECT* f = CONTAINING_RECORD(obj, FILE_OBJECT, Header);
    if (f->Fs && f->FileCtx) f->Fs->Close(f->Fs, f->FileCtx);
}

// ---------------------------------------------------------------------------
// IoOpenFile
// ---------------------------------------------------------------------------
__declspec(dllexport) KSTATUS IoOpenFile(const char* path, FILE_OBJECT **out) {
    if (!path || !out) return STATUS_INVALID_PARAMETER;

    char relPath[CWD_MAX];
    uint8_t drive = ResolvePath(path, relPath, sizeof(relPath));

    IO_FILESYSTEM* fs = gVolumes[drive].Fs;
    if (!fs) return STATUS_NOT_FOUND;

    void* fileCtx;
    KSTATUS s = fs->Open(fs, relPath, &fileCtx);
    if (!KSUCCESS(s)) return s;

    FILE_OBJECT* f = (FILE_OBJECT *)ObAllocateObject(ObjectTypeFile /* reuse slot */,
                                                      sizeof(FILE_OBJECT));
    if (!f) { fs->Close(fs, fileCtx); return STATUS_NO_MEMORY; }
    f->Header.Destroy = FileDestroy;
    f->Fs      = fs;
    f->FileCtx = fileCtx;
    f->Size    = fs->GetSize(fs, fileCtx);
    f->Offset  = 0;

    *out = f;
    return STATUS_SUCCESS;
}

// ---------------------------------------------------------------------------
// IoReadFile - sequential read at current offset
// ---------------------------------------------------------------------------
__declspec(dllexport) KSTATUS IoReadFile(FILE_OBJECT* file, void* buf, size_t size, size_t* bytesRead) {
    if (bytesRead) *bytesRead = 0;
    if (!file || !buf || !bytesRead) return STATUS_INVALID_PARAMETER;
    if (file->Offset >= file->Size || size == 0) return STATUS_SUCCESS;

    // Clamp to remaining bytes.
    size_t remaining = (size_t)(file->Size - file->Offset);
    if (size > remaining) size = remaining;

    KSTATUS s = file->Fs->Read(file->Fs, file->FileCtx,
                                file->Offset, buf, size, bytesRead);
    if (KSUCCESS(s)) file->Offset += *bytesRead;
    return s;
}

// ---------------------------------------------------------------------------
// IoWriteFile - sequential write at current offset
// ---------------------------------------------------------------------------
__declspec(dllexport) KSTATUS IoWriteFile(FILE_OBJECT* file, const void* buf,
                                          size_t size, size_t* bytesWritten) {
    if (!file || !buf || !bytesWritten) return STATUS_INVALID_PARAMETER;
    if (!file->Fs->Write) return STATUS_NOT_SUPPORTED;

    size_t written = 0;
    KSTATUS s = file->Fs->Write(file->Fs, file->FileCtx,
                                file->Offset, buf, size, &written);
    if (KSUCCESS(s)) {
        file->Offset += written;
        // Keep the cached size in sync if the file grew.
        if (file->Offset > file->Size)
            file->Size = file->Offset;
    }
    *bytesWritten = written;
    return s;
}

// ---------------------------------------------------------------------------
// IoCreateFile - create a new file and return it open for writing
// ---------------------------------------------------------------------------
__declspec(dllexport) KSTATUS IoCreateFile(const char* path, FILE_OBJECT** out) {
    if (!path || !out) return STATUS_INVALID_PARAMETER;

    char relPath[CWD_MAX];
    uint8_t drive = ResolvePath(path, relPath, sizeof(relPath));

    IO_FILESYSTEM* fs = gVolumes[drive].Fs;
    if (!fs)          return STATUS_NOT_FOUND;
    if (!fs->Create)  return STATUS_NOT_SUPPORTED;

    void* fileCtx;
    KSTATUS s = fs->Create(fs, relPath, &fileCtx);
    if (!KSUCCESS(s)) return s;

    FILE_OBJECT* f = (FILE_OBJECT*)ObAllocateObject(ObjectTypeFile,
                                                     sizeof(FILE_OBJECT));
    if (!f) { fs->Close(fs, fileCtx); return STATUS_NO_MEMORY; }
    f->Header.Destroy = FileDestroy;
    f->Fs      = fs;
    f->FileCtx = fileCtx;
    f->Size    = 0;
    f->Offset  = 0;

    *out = f;
    return STATUS_SUCCESS;
}

// ---------------------------------------------------------------------------
// IoDeleteFile - delete a file by path (does not require an open handle)
// ---------------------------------------------------------------------------
__declspec(dllexport) KSTATUS IoDeleteFile(const char* path) {
    if (!path) return STATUS_INVALID_PARAMETER;

    char relPath[CWD_MAX];
    uint8_t drive = ResolvePath(path, relPath, sizeof(relPath));

    IO_FILESYSTEM* fs = gVolumes[drive].Fs;
    if (!fs)         return STATUS_NOT_FOUND;
    if (!fs->Delete) return STATUS_NOT_SUPPORTED;

    return fs->Delete(fs, relPath);
}

__declspec(dllexport) KSTATUS IoRenameFile(const char* oldPath, const char* newPath) {
    if (!oldPath || !newPath) return STATUS_INVALID_PARAMETER;

    char oldRel[CWD_MAX];
    char newRel[CWD_MAX];
    uint8_t oldDrive = ResolvePath(oldPath, oldRel, sizeof(oldRel));
    uint8_t newDrive = ResolvePath(newPath, newRel, sizeof(newRel));
    if (oldDrive != newDrive) return STATUS_NOT_SUPPORTED;

    IO_FILESYSTEM* fs = gVolumes[oldDrive].Fs;
    if (!fs)          return STATUS_NOT_FOUND;
    if (!fs->Rename)  return STATUS_NOT_SUPPORTED;

    return fs->Rename(fs, oldRel, newRel);
}

__declspec(dllexport) KSTATUS IoCreateDir(const char* path) {
    if (!path) return STATUS_INVALID_PARAMETER;

    char relPath[CWD_MAX];
    uint8_t drive = ResolvePath(path, relPath, sizeof(relPath));

    IO_FILESYSTEM* fs = gVolumes[drive].Fs;
    if (!fs)           return STATUS_NOT_FOUND;
    if (!fs->CreateDir) return STATUS_NOT_SUPPORTED;

    return fs->CreateDir(fs, relPath);
}

__declspec(dllexport) KSTATUS IoRemoveDir(const char* path) {
    if (!path) return STATUS_INVALID_PARAMETER;

    char relPath[CWD_MAX];
    uint8_t drive = ResolvePath(path, relPath, sizeof(relPath));

    IO_FILESYSTEM* fs = gVolumes[drive].Fs;
    if (!fs)           return STATUS_NOT_FOUND;
    if (!fs->RemoveDir) return STATUS_NOT_SUPPORTED;

    return fs->RemoveDir(fs, relPath);
}

// ---------------------------------------------------------------------------
// IoTruncateFile - resize the file; updates the cached Size on success
// ---------------------------------------------------------------------------
__declspec(dllexport) KSTATUS IoTruncateFile(FILE_OBJECT* file, uint64_t newSize) {
    if (!file) return STATUS_INVALID_PARAMETER;
    if (!file->Fs->Truncate) return STATUS_NOT_SUPPORTED;

    KSTATUS s = file->Fs->Truncate(file->Fs, file->FileCtx, newSize);
    if (KSUCCESS(s)) {
        file->Size = newSize;
        // Pull the file pointer back if it now sits past EOF.
        if (file->Offset > file->Size)
            file->Offset = file->Size;
    }
    return s;
}

// ---------------------------------------------------------------------------
__declspec(dllexport) KSTATUS IoSeekFile(FILE_OBJECT* file, uint64_t offset) {
    if (!file) return STATUS_INVALID_PARAMETER;
    file->Offset = offset;
    return STATUS_SUCCESS;
}

__declspec(dllexport) uint64_t IoGetFileSize(FILE_OBJECT* file) {
    return file ? file->Size : 0;
}

__declspec(dllexport) void IoCloseFile(FILE_OBJECT* file) {
    if (file) ObDereferenceObject(&file->Header);
}

// ---------------------------------------------------------------------------
// Directory enumeration
// ---------------------------------------------------------------------------

static void DirDestroy(OBJECT* obj) {
    DIR_OBJECT* d = (DIR_OBJECT*)obj;  // Header is first member
    if (d->Fs && d->DirCtx) d->Fs->CloseDir(d->Fs, d->DirCtx);
}

__declspec(dllexport) KSTATUS IoOpenDir(const char* dirPath, DIR_OBJECT** out) {
    if (!out) return STATUS_INVALID_PARAMETER;

    // NULL or empty path means the CWD itself
    char relPath[CWD_MAX];
    uint8_t drive;

    if (!dirPath || dirPath[0] == '\0') {
        PROCESS* proc = KeGetCurrentProcess();
        if (!proc) return STATUS_ACCESS_DENIED;

        uintptr_t flags = KeAcquireSpinlock(&proc->CwdLock);
        memcpy(relPath, proc->CurrentDirectory, strlen(proc->CurrentDirectory) + 1);
        KeReleaseSpinlock(&proc->CwdLock, flags);
        
        // Note: For now, this assumes drive 0. If you add multi-drive support 
        // later, you should add a CurrentDrive field to the PROCESS struct.
        drive = 0; 
    } else {
        drive = ResolvePath(dirPath, relPath, sizeof(relPath));
    }

    IO_FILESYSTEM* fs = gVolumes[drive].Fs;
    if (!fs)          return STATUS_NOT_FOUND;
    if (!fs->OpenDir) return STATUS_NOT_SUPPORTED;

    void* dirCtx = NULL;
    KSTATUS s = fs->OpenDir(fs, relPath, &dirCtx);
    if (!KSUCCESS(s)) return s;

    DIR_OBJECT* d = (DIR_OBJECT*)ObAllocateObject(ObjectTypeDir, sizeof(DIR_OBJECT));
    if (!d) { fs->CloseDir(fs, dirCtx); return STATUS_NO_MEMORY; }
    
    d->Header.Destroy = DirDestroy;
    d->Fs     = fs;
    d->DirCtx = dirCtx;

    *out = d;
    return STATUS_SUCCESS;
}

__declspec(dllexport) KSTATUS IoReadDir(DIR_OBJECT* dir, KDIR_ENTRY* entryOut) {
    if (!dir || !entryOut) return STATUS_INVALID_PARAMETER;
    if (!dir->Fs->ReadDir) return STATUS_NOT_SUPPORTED;
    return dir->Fs->ReadDir(dir->Fs, dir->DirCtx, entryOut);
}

__declspec(dllexport) void IoCloseDir(DIR_OBJECT* dir) {
    if (dir) ObDereferenceObject(&dir->Header);
}