// Synthetic regression checks. Linux linker wrappers observe storage immediately
// before release and inject mprotect failures; they never enter production builds.
#include <algorithm>
#include <cerrno>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <vector>
#include "randomx.h"
#include "virtual_machine.hpp"
#include "allocator.hpp"
#include "virtual_memory.h"
#include "blake2/blake2.h"
#ifdef __linux__
#include <csignal>
#include <sys/mman.h>
#include <sys/resource.h>
#include <sys/wait.h>
#include <unistd.h>
#endif

static void require(bool ok, const char* message) {
    if (!ok) { std::fprintf(stderr, "FAIL: %s\n", message); std::exit(1); }
}
static bool zero(const void* p, size_t n) {
    const auto* bytes = static_cast<const unsigned char*>(p);
    return std::all_of(bytes, bytes + n, [](unsigned char b) { return b == 0; });
}

#ifdef __linux__
static bool track_jit = false;
static std::vector<void*> jit_buffers;
static void* scratch_to_check = nullptr;
static unsigned jit_wipes = 0, scratch_wipes = 0;
static int fail_protection = -1;
extern "C" void* __real_allocMemoryPages(size_t);
extern "C" void __real_freePagedMemory(void*, size_t);
extern "C" void __real_free(void*);
extern "C" int __real_mprotect(void*, size_t, int);
extern "C" void* __wrap_allocMemoryPages(size_t n) {
    void* p = __real_allocMemoryPages(n);
    if (track_jit && p) jit_buffers.push_back(p);
    return p;
}
extern "C" void __wrap_freePagedMemory(void* p, size_t n) {
    auto i = std::find(jit_buffers.begin(), jit_buffers.end(), p);
    if (i != jit_buffers.end()) {
        require(zero(p, n), "JIT allocation was not zero before unmapping");
        jit_buffers.erase(i);
        ++jit_wipes;
    }
    __real_freePagedMemory(p, n);
}
extern "C" void __wrap_free(void* p) {
    if (p && p == scratch_to_check) {
        require(zero(p, randomx::ScratchpadSize), "scratchpad was not zero before free");
        scratch_to_check = nullptr;
        ++scratch_wipes;
    }
    __real_free(p);
}
extern "C" int __wrap_mprotect(void* p, size_t n, int protection) {
    if (protection == fail_protection) { errno = EPERM; return -1; }
    return __real_mprotect(p, n, protection);
}
static void expect_abort(randomx_vm* vm, int protection, bool hash) {
    const pid_t child = fork();
    require(child >= 0, "fork");
    if (child == 0) {
        rlimit limit{0, 0}; setrlimit(RLIMIT_CORE, &limit);
        fail_protection = protection;
        if (hash) {
            unsigned char out[32];
            randomx_calculate_hash(vm, "failure injection", 17, out);
        } else {
            void* p = allocMemoryPages(4096);
            if (protection == (PROT_READ | PROT_WRITE)) setPagesRW(p, 4096);
            else if (protection == (PROT_READ | PROT_EXEC)) setPagesRX(p, 4096);
            else setPagesRWX(p, 4096);
        }
        _exit(88); // A failed permission transition must never return normally.
    }
    int status;
    require(waitpid(child, &status, 0) == child, "waitpid");
    require(WIFSIGNALED(status) && WTERMSIG(status) == SIGABRT, "permission failure did not abort");
}
#endif

int main(int argc, char**) {
    const bool full = argc > 1;
    const auto flags = randomx_get_flags();
    require(flags & RANDOMX_FLAG_HARD_AES, "hardware AES required for this fixture");
    require(flags & RANDOMX_FLAG_JIT, "native JIT required for this fixture");
    auto* cache = randomx_alloc_cache(flags);
    require(cache != nullptr, "cache allocation");
    randomx_init_cache(cache, "test key 000", 12);
    auto* dataset = full ? randomx_alloc_dataset(RANDOMX_FLAG_DEFAULT) : nullptr;
    require(!full || dataset, "dataset allocation");
    if (full) randomx_init_dataset(dataset, cache, 0, randomx_dataset_item_count());
    const auto memory = full ? RANDOMX_FLAG_FULL_MEM : RANDOMX_FLAG_DEFAULT;
    auto* reference = randomx_create_vm(RANDOMX_FLAG_V2 | RANDOMX_FLAG_HARD_AES | memory, cache, dataset);
    require(reference != nullptr, "interpreted VM allocation");
#ifdef __linux__
    track_jit = true;
#endif
    auto* vm = randomx_create_vm(flags | RANDOMX_FLAG_V2 | RANDOMX_FLAG_SECURE | memory, cache, dataset);
    require(vm != nullptr, "secure JIT VM allocation");
#ifdef __linux__
    track_jit = false;
    for (int prot : {PROT_READ | PROT_WRITE, PROT_READ | PROT_EXEC, PROT_READ | PROT_WRITE | PROT_EXEC})
        expect_abort(vm, prot, false);
    expect_abort(vm, PROT_READ | PROT_WRITE, true);
    expect_abort(vm, PROT_READ | PROT_EXEC, true);
    std::puts("5 injected permission failures aborted before returning a hash");
#endif
    unsigned char output[32], expected[32];
    for (unsigned i = 0; i < 16; ++i) {
        unsigned char input[80];
        for (unsigned j = 0; j < sizeof(input); ++j) input[j] = static_cast<unsigned char>(i * 31 + j);
        randomx_calculate_hash(vm, input, sizeof(input), output);
        randomx_calculate_hash(reference, input, sizeof(input), expected);
        require(std::memcmp(output, expected, 32) == 0, "JIT/interpreter differential");
    }
    std::printf("16 secure-JIT/interpreted %s hashes matched\n", full ? "full" : "light");

    for (auto* machine : {vm, reference}) {
        // Capture byte addresses while the object is live, then inspect its
        // retained allocation after the virtual destructor and before free.
        auto* registers = reinterpret_cast<unsigned char*>(machine->getRegisterFile());
        auto* program = reinterpret_cast<const unsigned char*>(&machine->getProgram());
        auto* temporary = reinterpret_cast<unsigned char*>(machine->tempHash);
        require(!zero(registers, sizeof(randomx::RegisterFile)), "fixture registers unexpectedly zero");
#ifdef __linux__
        scratch_to_check = const_cast<void*>(machine->getScratchpad());
#endif
        machine->~randomx_vm();
        require(zero(registers, sizeof(randomx::RegisterFile)), "register remanence after destructor");
        require(zero(program, sizeof(randomx::Program)), "program remanence after destructor");
        require(zero(temporary, 64), "temporary hash remanence after destructor");
        blake2b(expected, sizeof(expected), registers, sizeof(randomx::RegisterFile), nullptr, 0);
        require(std::memcmp(output, expected, 32) != 0, "last hash recoverable after destructor");
        randomx::AlignedAllocator<randomx::CacheLineSize>::freeMemory(machine, 0);
    }
    std::puts("JIT and interpreted VM register/program/tempHash storage erased; final output not recoverable");
#ifdef __linux__
    require(scratch_wipes == 2 && jit_wipes == 1 && jit_buffers.empty(), "missing deallocation observations");
    std::puts("Both scratchpads and the complete JIT allocation were zero before release");
#endif
    if (dataset) randomx_release_dataset(dataset);
    randomx_release_cache(cache);
    return 0;
}