// Shared writable mmap regression test.
// Verifies that writable mapped-file views dirty the coherent file cache and flush through the normal file path.

#include <stdint.h>
#include <stddef.h>
#include <string.h>
#include <ntlite/syscall.h>
#include <ntlite/path.h>

#define PAGE_SIZE_BYTES 4096u
#define TEST_SIZE       (PAGE_SIZE_BYTES * 3u)

static uint8_t gInitial[TEST_SIZE];
static uint8_t gExpected[TEST_SIZE];
static uint8_t gReadback[TEST_SIZE];

static void print(const char* s) {
    SYS_WRITE_CONSOLE_REQUEST req = { sizeof(SYS_WRITE_CONSOLE_REQUEST), SysSvFromCstr(s) };
    SysWriteConsole(&req);
}

static void print_u64(uint64_t v) {
    char buf[32];
    int i = 0;
    if (v == 0) { print("0"); return; }
    while (v && i < (int)sizeof(buf)) {
        buf[i++] = (char)('0' + (v % 10));
        v /= 10;
    }
    while (i--) {
        char c[2] = { buf[i], 0 };
        print(c);
    }
}

static void print_hex_uintptr(uintptr_t value) {
    static const char digits[] = "0123456789abcdef";
    char text[2 + sizeof(uintptr_t) * 2 + 1];
    size_t pos = 0;
    text[pos++] = '0';
    text[pos++] = 'x';
    for (int shift = (int)(sizeof(uintptr_t) * 8u - 4u); shift >= 0; shift -= 4)
        text[pos++] = digits[(value >> (uint32_t)shift) & 0xFu];
    text[pos] = 0;
    print(text);
}

// Prints the VM region containing one address without requiring physical-address disclosure.
static void print_vm_region(const char* name, const void* address) {
    SYS_VM_REGION_INFO info;
    memset(&info, 0, sizeof(info));
    KSTATUS status = SysQueryVm((void*)address, &info);

    print("shared_mmap: vm ");
    print(name);
    print(" address=");
    print_hex_uintptr((uintptr_t)address);
    if (status != NTLSTATUS_SUCCESS) {
        print(" query-status=");
        print_u64(status);
        print("\n");
        return;
    }

    print(" base=");
    print_hex_uintptr((uintptr_t)info.Base);
    print(" pages=");
    print_u64(info.PageCount);
    print(" type=");
    print_u64(info.Type);
    print("\n");
}

static void exit_now(int code) {
    while (true) SysExitProcess(code);
}

static void fail(const char* msg) {
    print("shared_mmap: ");
    print(msg);
    print("\n");
    exit_now(1);
}

static void fail_status(const char* what, KSTATUS status) {
    print("shared_mmap: ");
    print(what);
    print(" failed status=");
    print_u64(status);
    print("\n");
    exit_now(1);
}

// Reports the first byte mismatch before terminating the regression test.
static void fail_mismatch(const char* what, const uint8_t* actual, const uint8_t* expected, size_t size) {
    for (size_t i = 0; i < size; i++) {
        uint8_t actualByte = actual[i];
        uint8_t expectedByte = expected[i];
        if (actualByte == expectedByte) continue;
        print("shared_mmap: ");
        print(what);
        print(" offset=");
        print_u64(i);
        print(" expected=");
        print_u64(expectedByte);
        print(" actual=");
        print_u64(actualByte);
        print("\n");
        exit_now(1);
    }
    fail(what);
}

static uint32_t query_cache(SYS_FILE_CACHE_STATS* out) {
    memset(out, 0, sizeof(*out));
    return SysQueryKernelStats(SYS_KERNEL_STATS_FILE_CACHE, out, sizeof(*out), 0);
}

static void delete_path(const char* path) {
    (void)NtlDeleteFile(SysSvFromCstr(path));
}

static KSTATUS create_temp(uint32_t* fileOut, const char** pathOut) {
    static const char* paths[] = {
        "shared_mmap.tmp",
    };

    for (size_t i = 0; i < sizeof(paths) / sizeof(paths[0]); ++i) {
        uint32_t h = 0;
        delete_path(paths[i]);
        KSTATUS s = NtlOpenFile(SysSvFromCstr(paths[i]),
                                         SYS_FILE_ACCESS_READ | SYS_FILE_ACCESS_WRITE | SYS_FILE_ACCESS_QUERY,
                                         SYS_FILE_CREATE_NEW,
                                         SYS_FILE_FLAG_NONE,
                                         &h);
        if (s == NTLSTATUS_SUCCESS) {
            *fileOut = h;
            if (pathOut) *pathOut = paths[i];
            return NTLSTATUS_SUCCESS;
        }
    }
    return NTLSTATUS_NOT_FOUND;
}

static KSTATUS open_path(const char* path, uint32_t* fileOut) {
    uint32_t h = 0;
    KSTATUS s = NtlOpenFile(SysSvFromCstr(path),
                                     SYS_FILE_ACCESS_READ | SYS_FILE_ACCESS_WRITE | SYS_FILE_ACCESS_QUERY,
                                     SYS_FILE_OPEN_EXISTING,
                                     SYS_FILE_FLAG_NONE,
                                     &h);
    if (s == NTLSTATUS_SUCCESS)
        *fileOut = h;
    return s;
}

static void checked_seek(uint32_t file, uint64_t offset, const char* what) {
    KSTATUS s = SysSeekFile(file, offset);
    if (s != NTLSTATUS_SUCCESS) fail_status(what, s);
}

static void checked_write(uint32_t file, const void* buffer, size_t size, const char* what) {
    size_t done = 0;
    SYS_FILE_IO_REQUEST req = SysMakeFileIoRequest(file, buffer, size, &done);
    KSTATUS s = SysWriteFile(&req);
    if (s != NTLSTATUS_SUCCESS) fail_status(what, s);
    if (done != size) fail("short write");
}

static void checked_read(uint32_t file, void* buffer, size_t size, size_t* gotOut, const char* what) {
    size_t got = 0;
    memset(buffer, 0xCC, size);
    SYS_FILE_IO_REQUEST req = SysMakeFileIoRequest(file, buffer, size, &got);
    KSTATUS s = SysReadFile(&req);
    if (s != NTLSTATUS_SUCCESS) fail_status(what, s);
    if (gotOut) *gotOut = got;
}

// Returns the deterministic initial byte used to seed the mapped file.
static uint8_t initial_byte(size_t index) {
    return (uint8_t)(0x20u + (uint8_t)(index * 7u + (index >> 5)));
}

// Returns the expected byte after the first shared-mapping write range.
static uint8_t expected_after_first_mapped_write(size_t index) {
    if (index >= 128u && index < 128u + 3000u)
        return (uint8_t)(0xA0u + (uint8_t)index);
    return initial_byte(index);
}

// Returns the expected byte after both shared-mapping write ranges.
static uint8_t expected_after_mapped_writes(size_t index) {
    if (index >= 128u && index < 128u + 3000u)
        return (uint8_t)(0xA0u + (uint8_t)index);
    if (index >= PAGE_SIZE_BYTES + 17u && index < PAGE_SIZE_BYTES + 17u + 2048u)
        return (uint8_t)(0x50u + (uint8_t)(index >> 1));
    return initial_byte(index);
}

// Verifies the mutable expected buffer against the deterministic write pattern.
static void verify_expected_buffer(const char* stage, int secondRangeWritten) {
    for (size_t i = 0; i < TEST_SIZE; ++i) {
        uint8_t actualByte = gExpected[i];
        uint8_t expectedByte = secondRangeWritten
            ? expected_after_mapped_writes(i)
            : expected_after_first_mapped_write(i);
        if (actualByte == expectedByte)
            continue;

        print("shared_mmap: expected buffer corrupted ");
        print(stage);
        print(" offset=");
        print_u64(i);
        print(" expected=");
        print_u64(expectedByte);
        print(" actual=");
        print_u64(actualByte);
        print("\n");
        exit_now(1);
    }
}

// Fills the initial file data and its comparison buffer.
static void fill_initial(void) {
    for (size_t i = 0; i < TEST_SIZE; ++i) {
        gInitial[i] = initial_byte(i);
        gExpected[i] = gInitial[i];
    }
}

int main(void) {
    print("shared_mmap: start\n");
    fill_initial();

    uint32_t file = 0;
    const char* path = 0;
    uint32_t s = create_temp(&file, &path);
    if (s != NTLSTATUS_SUCCESS) fail_status("create temp", s);

    print("shared_mmap: creating temporary file ");
    print(path ? path : "<unknown>");
    print("\n");

    checked_write(file, gInitial, TEST_SIZE, "initial write");
    s = SysFlushFile(file);
    if (s != NTLSTATUS_SUCCESS) fail_status("initial flush", s);

    uint32_t mapping = 0;
    SYS_CREATE_FILE_MAPPING_REQUEST createMapping = {
        sizeof(SYS_CREATE_FILE_MAPPING_REQUEST), file, SYS_VM_READWRITE,
        SYS_MAP_FLAG_NONE, TEST_SIZE, &mapping
    };
    s = SysCreateFileMapping(&createMapping);
    if (s != NTLSTATUS_SUCCESS) fail_status("create mapping", s);

    uint8_t* view = 0;
    SYS_MAP_FILE_REQUEST mapFile = {
        sizeof(SYS_MAP_FILE_REQUEST), mapping,
        SYS_MAP_ACCESS_READ | SYS_MAP_ACCESS_WRITE, SYS_MAP_FLAG_NONE,
        0, TEST_SIZE, (void**)&view
    };
    s = SysMapFile(&mapFile);
    if (s != NTLSTATUS_SUCCESS) fail_status("map shared write", s);

    print_vm_region("view", view);
    print_vm_region("expected", gExpected);
    print_vm_region("readback", gReadback);

    // Keep the source value live across the faulting mapped store. This also
    // verifies that a repaired page fault preserves the interrupted SIMD/GPR state.
    for (size_t i = 128; i < 128 + 3000; ++i) {
        uint8_t value = (uint8_t)(0xA0u + (uint8_t)i);
        view[i] = value;
        gExpected[i] = value;
    }
    verify_expected_buffer("after first mapped write range", 0);

    for (size_t i = PAGE_SIZE_BYTES + 17; i < PAGE_SIZE_BYTES + 17 + 2048; ++i) {
        uint8_t value = (uint8_t)(0x50u + (uint8_t)(i >> 1));
        view[i] = value;
        gExpected[i] = value;
    }
    verify_expected_buffer("after mapped writes", 1);

    SYS_FILE_CACHE_STATS beforeFlush;
    SYS_FILE_CACHE_STATS afterFlush;
    if (query_cache(&beforeFlush) != NTLSTATUS_SUCCESS) fail("pre-flush stats failed");
    verify_expected_buffer("after pre-flush stats", 1);

    s = SysFlushFile(file);
    if (s != NTLSTATUS_SUCCESS) fail_status("flush mapped writes", s);
    verify_expected_buffer("after mapped flush", 1);

    if (query_cache(&afterFlush) != NTLSTATUS_SUCCESS) fail("post-flush stats failed");
    verify_expected_buffer("after post-flush stats", 1);
    if (afterFlush.MappedDirtyPages <= beforeFlush.MappedDirtyPages)
        fail("mapped dirty pages were not harvested");
    if (afterFlush.WritebackBytes <= beforeFlush.WritebackBytes)
        fail("mapped writes did not write back");

    checked_seek(file, 0, "seek after mapped flush");
    size_t got = 0;
    checked_read(file, gReadback, TEST_SIZE, &got, "read after mapped flush");
    verify_expected_buffer("after mapped-flush read", 1);
    if (got != TEST_SIZE) fail("read after mapped flush size mismatch");
    if (memcmp(gReadback, gExpected, TEST_SIZE) != 0)
        fail_mismatch("mapped flush data mismatch", gReadback, gExpected, TEST_SIZE);

    for (size_t i = PAGE_SIZE_BYTES * 2 + 3; i < PAGE_SIZE_BYTES * 2 + 900; ++i) {
        uint8_t value = (uint8_t)(0xC0u ^ (uint8_t)i);
        view[i] = value;
        gExpected[i] = value;
    }

    s = SysFreeVm(view);
    view = 0;
    if (s != NTLSTATUS_SUCCESS) fail_status("unmap", s);

    s = SysCloseHandle(mapping);
    mapping = 0;
    if (s != NTLSTATUS_SUCCESS) fail_status("close mapping", s);

    s = SysCloseHandle(file);
    file = 0;
    if (s != NTLSTATUS_SUCCESS) fail_status("close file", s);

    s = open_path(path, &file);
    if (s != NTLSTATUS_SUCCESS) fail_status("reopen", s);

    checked_read(file, gReadback, TEST_SIZE, &got, "read after close");
    if (got != TEST_SIZE) fail("read after close size mismatch");
    if (memcmp(gReadback, gExpected, TEST_SIZE) != 0)
        fail_mismatch("unmap/close data mismatch", gReadback, gExpected, TEST_SIZE);

    s = SysCloseHandle(file);
    file = 0;
    if (s != NTLSTATUS_SUCCESS) fail_status("close reopened", s);
    delete_path(path);

    print("shared_mmap: mapped-dirty-pages+=");
    print_u64(afterFlush.MappedDirtyPages - beforeFlush.MappedDirtyPages);
    print(" writeback-bytes+=");
    print_u64(afterFlush.WritebackBytes - beforeFlush.WritebackBytes);
    print("\n");
    print("shared_mmap: PASS\n");
    exit_now(0);
    return 0;
}