#include "shim/trace/selftest.h"

#include <stdexcept>

namespace shim::selftest {

using trace::Tv;

std::uint32_t SHIM_CDECL Fill(std::uint8_t* buf, std::uint32_t n, std::uint32_t seed) {
    std::uint32_t x = seed, sum = 0;
    for (std::uint32_t i = 0; i < n; ++i) {
        x = x * 1664525u + 1013904223u;
        buf[i] = static_cast<std::uint8_t>(x >> 24);
        sum = sum * 31u + buf[i];
    }
    return sum;
}

// Same contract, written differently (rolling checksum accumulated after the fill).
std::uint32_t SHIM_CDECL FillOurs(std::uint8_t* buf, std::uint32_t n, std::uint32_t seed) {
    std::uint32_t x = seed;
    for (std::uint32_t i = 0; i < n; ++i) {
        x = 1664525u * x + 1013904223u;
        buf[i] = static_cast<std::uint8_t>((x >> 24) & 0xffu);
    }
    std::uint32_t sum = 0;
    for (std::uint32_t i = 0; i < n; ++i) sum = sum * 31u + buf[i];
    return sum;
}

std::uint32_t SHIM_CDECL FillWrong(std::uint8_t* buf, std::uint32_t n, std::uint32_t seed) {
    const std::uint32_t sum = FillOurs(buf, n, seed);
    if (n) buf[n - 1] = static_cast<std::uint8_t>(buf[n - 1] + 1);
    return sum;  // checksum still reports the original's answer: only the region diverges
}

std::uint32_t SHIM_CDECL FillThrows(std::uint8_t*, std::uint32_t, std::uint32_t) {
    throw std::runtime_error("selftest: deliberate throw");
}

// Fills the blob's buffer AND bumps a counter the hook's Result region does not cover.
std::uint32_t SHIM_CDECL FillCounted(Blob* b, std::uint32_t n, std::uint32_t seed) {
    if (!b) return 0;
    ++b->calls;
    return Fill(b->buf, n, seed);
}

// The faithful-looking reimplementation: correct buffer, no counter. This is B3 in miniature.
std::uint32_t SHIM_CDECL FillCountedOurs(Blob* b, std::uint32_t n, std::uint32_t seed) {
    if (!b) return 0;
    return FillOurs(b->buf, n, seed);
}

void FillHook::describe_args(std::vector<Tv>& out, std::uint8_t* buf, std::uint32_t n, std::uint32_t seed) {
    out.push_back(trace::tv::ptr(buf).named("buf"));
    out.push_back(trace::tv::u32(n).named("n"));
    out.push_back(trace::tv::u32(seed).named("seed"));
}

Tv FillHook::describe_ret(std::uint32_t r) { return trace::tv::u32(r); }

void FillHook::regions(std::vector<trace::Region>& out, std::uint8_t* buf, std::uint32_t n, std::uint32_t) {
    trace::Region r;
    r.name = "buf";
    r.ptr = buf;
    r.size = n;
    out.push_back(r);
}

FillHook::Args FillHook::rebind(trace::Scratch& s, std::uint8_t*, std::uint32_t n, std::uint32_t seed) {
    return Args(s.as<std::uint8_t>(0), n, seed);
}

void FillGuardHook::describe_args(std::vector<Tv>& out, Blob* b, std::uint32_t n, std::uint32_t seed) {
    out.push_back(trace::tv::ptr(b).named("blob"));
    out.push_back(trace::tv::u32(n).named("n"));
    out.push_back(trace::tv::u32(seed).named("seed"));
}

Tv FillGuardHook::describe_ret(std::uint32_t r) { return trace::tv::u32(r); }

void FillGuardHook::regions(std::vector<trace::Region>& out, Blob* b, std::uint32_t n, std::uint32_t) {
    trace::Region r;
    r.name = "buf";
    r.ptr = b ? b->buf : nullptr;
    r.size = n;
    out.push_back(r);

    // The whole object. Nothing our side writes lands here; anything the ORIGINAL moves outside
    // `buf` comes back as an undeclared write.
    trace::Region g;
    g.name = "blob";
    g.ptr = b;
    g.size = sizeof(Blob);
    g.kind = trace::Region::Kind::Guard;
    out.push_back(g);
}

FillGuardHook::Args FillGuardHook::rebind(trace::Scratch& s, Blob*, std::uint32_t n, std::uint32_t seed) {
    // Scratch holds Result regions only, so index 0 is `buf` even though a guard was declared.
    // `buf` is at offset 0 of a Blob and `ours` writes nothing else, so the copy stands in for
    // the object.
    return Args(reinterpret_cast<Blob*>(s.ptr(0)), n, seed);
}

std::uint32_t run_once(trace::Mode mode) {
    using H = trace::Hook<FillHook>;
    H::original = &Fill;
    H::mode = mode;
    std::uint8_t buf[64] = {};
    return H::detour()(buf, sizeof buf, 0x5eed);
}

}  // namespace shim::selftest