// System root volume selection. This file retains the structurally selected mounted volume for bootstrap capability-relative I/O.

#include <stdbool.h>
#include <stddef.h>
#include <stdint.h>
#include <kernel/ke_system_root.h>
#include <kernel/ke_debug.h>
#include <kernel/io_volume.h>
#include <kernel/io_partition.h>
#include <kernel/ke_sched.h>

static IO_VOLUME* gSystemRootVolume;
static bool gSystemRootReady;

#define KE_SYSTEM_ROOT_DISCOVERY_TIMEOUT_MS 30000ull
#define KE_SYSTEM_ROOT_WAIT_SLICE_MS         1000ull
#define KE_SYSTEM_ROOT_LOG_INTERVAL_MS       5000ull

// Returns the current scheduler time in milliseconds.
static uint64_t KiSystemRootNowMilliseconds(void) {
    KE_SCHED_STATS stats;
    KeGetSchedulerStats(&stats);
    return stats.Ticks * SCHED_TICK_MS;
}

// Resolves the unique configured partition identity to its mounted filesystem volume.
static KSTATUS KiTrySelectSystemRootVolume(uint64_t* mountedOut) {
    IO_VOLUME_STATS stats;
    IoQueryVolumeStats(&stats);
    if (mountedOut)
        *mountedOut = stats.MountedVolumes;

    IO_PARTITION_INFO systemPartition;
    KSTATUS status = IoQuerySystemPartition(&systemPartition);
    if (!KSUCCESS(status))
        return status;

    IO_VOLUME* volume = (IO_VOLUME*)0;
    status = IoReferenceVolumeByBacking(systemPartition.DiskIndex,
                                        systemPartition.StartLba,
                                        &volume);
    if (!KSUCCESS(status))
        return status;

    gSystemRootVolume = volume; // Retain the mounted-volume reference for the system lifetime.
    gSystemRootReady = true;
    DbgPrint("System root: volume=%u disk=%u start_lba=%llu sectors=%llu\n",
             IoGetVolumeInstanceId(volume), systemPartition.DiskIndex,
             (unsigned long long)systemPartition.StartLba,
             (unsigned long long)systemPartition.SectorCount);
    return NTLSTATUS_SUCCESS;
}

// Waits for the configured structural system partition to be discovered and mounted.
KSTATUS KeInitializeSystemRootVolume(void) {
    if (gSystemRootReady)
        return NTLSTATUS_SUCCESS;

    uint64_t start = KiSystemRootNowMilliseconds();
    uint64_t deadline = start + KE_SYSTEM_ROOT_DISCOVERY_TIMEOUT_MS;
    uint64_t nextLog = start;

    for (;;) {
        uint64_t generation = IoGetVolumeGeneration();
        uint64_t mounted = 0;
        KSTATUS status = KiTrySelectSystemRootVolume(&mounted);
        if (KSUCCESS(status))
            return status;
        if (status == NTLSTATUS_OBJECT_NAME_COLLISION) {
            DbgPrint("System root: configured system_volume matches multiple mounted volumes\n");
            return status;
        }
        if (status != NTLSTATUS_NOT_FOUND)
            return status;

        uint64_t now = KiSystemRootNowMilliseconds();
        if (now >= deadline) {
            DbgPrint("System root: discovery timed out after %llu ms; mounted=%llu generation=%llu\n",
                     (unsigned long long)KE_SYSTEM_ROOT_DISCOVERY_TIMEOUT_MS,
                     (unsigned long long)mounted,
                     (unsigned long long)generation);
            return NTLSTATUS_TIMEOUT;
        }

        if (now >= nextLog) {
            DbgPrint("System root: waiting for configured partition; mounted=%llu generation=%llu remaining=%llu ms\n",
                     (unsigned long long)mounted,
                     (unsigned long long)generation,
                     (unsigned long long)(deadline - now));
            nextLog = now + KE_SYSTEM_ROOT_LOG_INTERVAL_MS;
        }

        uint64_t remaining = deadline - now;
        uint64_t wait = remaining < KE_SYSTEM_ROOT_WAIT_SLICE_MS
                      ? remaining
                      : KE_SYSTEM_ROOT_WAIT_SLICE_MS;
        KSTATUS waitStatus = IoWaitForVolumeChange(generation, wait);
        if (waitStatus != NTLSTATUS_SUCCESS && waitStatus != NTLSTATUS_TIMEOUT)
            return waitStatus;
    }
}

// Returns whether the system root volume has been selected.
bool KeSystemRootReady(void) {
    return gSystemRootReady;
}

// References the selected system volume for capability transfer or kernel-relative I/O.
KSTATUS KeReferenceSystemRootVolume(IO_VOLUME** volumeOut) {
    if (!volumeOut)
        return NTLSTATUS_INVALID_PARAMETER;
    *volumeOut = (IO_VOLUME*)0;
    if (!gSystemRootReady || !gSystemRootVolume)
        return NTLSTATUS_NOT_FOUND;
    if (!IoVolumeTryReferenceMounted(gSystemRootVolume))
        return NTLSTATUS_OBJECT_DEAD;
    *volumeOut = gSystemRootVolume;
    return NTLSTATUS_SUCCESS;
}

// Opens a canonical native path rooted directly at the selected system volume.
KSTATUS KeOpenSystemRootFile(const char* relative, FILE_OBJECT** out) {
    if (!relative || !out)
        return NTLSTATUS_INVALID_PARAMETER;
    *out = NULL;
    if (!gSystemRootReady || !gSystemRootVolume)
        return NTLSTATUS_NOT_FOUND;
    return IoOpenFileAtVolume(gSystemRootVolume, SysSvFromCstr(relative), out);
}

// Creates a canonical native path rooted directly at the selected system volume.
KSTATUS KeCreateSystemRootFile(const char* relative, FILE_OBJECT** out) {
    if (!relative || !out)
        return NTLSTATUS_INVALID_PARAMETER;
    *out = NULL;
    if (!gSystemRootReady || !gSystemRootVolume)
        return NTLSTATUS_NOT_FOUND;
    return IoCreateFileAtVolume(gSystemRootVolume, SysSvFromCstr(relative), out);
}

// Deletes a canonical native path rooted directly at the selected system volume.
KSTATUS KeDeleteSystemRootFile(const char* relative) {
    if (!relative)
        return NTLSTATUS_INVALID_PARAMETER;
    if (!gSystemRootReady || !gSystemRootVolume)
        return NTLSTATUS_NOT_FOUND;
    return IoDeleteFileAtRoot(&gSystemRootVolume->Header, SysSvFromCstr(relative));
}