#include "shim/hooks/probe_entry.h"

#include <cstdarg>
#include <cstdio>
#include <cstdlib>
#include <cstring>

#if defined(_WIN32)
#define WIN32_LEAN_AND_MEAN
#include <windows.h>

#include "MinHook.h"
#endif

#include "generated/sots_addresses.h"

// ---- the counters ------------------------------------------------------------------------------
//
// Plain 32-bit counters, incremented from a stub that has already saved the flags, so no atomic is
// needed for correctness of the arithmetic on x86 and none is used: the turn pipeline is
// single-threaded (lane Z §7) and a torn read here would cost a count, not a crash. If a probe ever
// reads a wild number, that is itself the finding -- it would mean the address is entered off the
// turn thread.

namespace {
constexpr std::size_t kMaxProbes = 16;
std::uint32_t g_calls[kMaxProbes];
std::uint32_t g_total[kMaxProbes];
}  // namespace

extern "C" void ProbeHit(int index) {
    if (index < 0 || static_cast<std::size_t>(index) >= kMaxProbes) return;
    ++g_calls[index];
    ++g_total[index];
}

// ---- one trampoline slot and one asm stub per probe ---------------------------------------------
//
// The stub is written out per index rather than generated at runtime because a runtime-generated
// thunk would need an executable allocation and a relocation, and MinHook already owns that
// problem. Sixteen slots is the cap; `kProbes` below is shorter and the rest are inert.

#if defined(_WIN32)
#define PROBE_STUB(i)                                    \
    extern "C" void* g_tr_probe##i;                      \
    void* g_tr_probe##i = nullptr;                       \
    extern "C" void ProbeStub##i(void);                  \
    asm(".text\n"                                        \
        ".globl _ProbeStub" #i "\n"                      \
        "_ProbeStub" #i ":\n"                            \
        "    pushfl\n"                                   \
        "    pushal\n"                                   \
        "    pushl $" #i "\n"                            \
        "    call _ProbeHit\n"                           \
        "    addl $4, %esp\n"                            \
        "    popal\n"                                    \
        "    popfl\n"                                    \
        "    jmp *_g_tr_probe" #i "\n")

PROBE_STUB(0);
PROBE_STUB(1);
PROBE_STUB(2);
PROBE_STUB(3);
PROBE_STUB(4);
PROBE_STUB(5);
PROBE_STUB(6);
PROBE_STUB(7);
PROBE_STUB(8);
PROBE_STUB(9);
PROBE_STUB(10);
PROBE_STUB(11);
#undef PROBE_STUB
#endif

namespace shim::hooks {

namespace {

namespace A = sots::addr;

struct ProbeDef {
    const char* name;
    std::uint32_t rva;
    void* stub;
    void** trampoline;
};

std::uintptr_t g_exe_base = 0;
void (*g_log)(const char*) = nullptr;
bool g_installed[kMaxProbes];

void logf(const char* fmt, ...) {
    if (!g_log) return;
    char line[512];
    va_list ap;
    va_start(ap, fmt);
    std::vsnprintf(line, sizeof line, fmt, ap);
    va_end(ap);
    g_log(line);
}

#if defined(_WIN32)
#define PROBE(name, rva, i) ProbeDef{name, rva, reinterpret_cast<void*>(&ProbeStub##i), &g_tr_probe##i}
#else
#define PROBE(name, rva, i) ProbeDef{name, rva, nullptr, nullptr}
#endif

// The probe set. Order is the report order and the stub index, so it is fixed.
//
// Rows 0-3 are the four phase-23/33 callees lane V2 proved can reach the generator; rows 4-6 are
// the three inner functions those bodies reach the draw *through*, so a row-0..3 hit with no
// row-4..6 hit localises the gate to the outer body. Rows 7-8 are the trade-raid pair (P4). Rows
// 9-11 are CONTROLS, and they are the reason a zero in rows 0-6 can be believed: 9 and 10 are two
// phase-23/33 callees lane V2 measured as draw-free, and 11 is `EncounterDetect_Run`, which every
// prior lane has observed running once per turn. If the controls read 0 the instrument is broken,
// not the game (method rule 1 -- a hook that compares nothing prints a clean verdict).
const ProbeDef kProbes[] = {
    PROBE("Game::ServerSpyManager::vslot13", A::ServerSpyManager_vslot13, 0),
    PROBE("Game::ServerSpyManager::vslot14", A::ServerSpyManager_vslot14, 1),
    PROBE("Game::ServerTradeManagerImpl::vslot13", A::ServerTradeManagerImpl_vslot13, 2),
    PROBE("Game::ServerTradeManagerImpl::vslot15", A::ServerTradeManagerImpl_vslot15, 3),
    PROBE("Game::SpyManager::Slot13RngCallee", A::SpyManager_Slot13RngCallee, 4),
    PROBE("Game::TradeManager::Slot13RngCalleeA", A::TradeManager_Slot13RngCalleeA, 5),
    PROBE("Game::TradeManager::Slot13RngCalleeB", A::TradeManager_Slot13RngCalleeB, 6),
    PROBE("Game::ServerTradeManager::CreateRaidEncounter", A::ServerTradeManager_CreateRaidEncounter, 7),
    PROBE("Game::ServerTradeManager::GenerateTradeRaidEncounters",
          A::ServerTradeManager_GenerateTradeRaidEncounters, 8),
    PROBE("Game::ServerSpyManager::vslot15 [control]", A::ServerSpyManager_vslot15, 9),
    PROBE("Game::ServerTradeManagerImpl::vslot14 [control]", A::ServerTradeManagerImpl_vslot14, 10),
    PROBE("Game::EncounterDetect::Run [control]", A::EncounterDetect_Run, 11),
};
#undef PROBE

constexpr std::size_t kProbeCount = sizeof kProbes / sizeof kProbes[0];
static_assert(kProbeCount <= kMaxProbes, "add more PROBE_STUB() slots");

}  // namespace

namespace {
std::size_t g_install_count = kProbeCount;  // `probes=` overrides; default is all of them
}  // namespace

bool probe_config(const char* key, const char* value, std::size_t* count_out) {
    if (std::strcmp(key, "probes") != 0) return false;
    if (std::strcmp(value, "off") == 0 || std::strcmp(value, "none") == 0) {
        g_install_count = 0;
    } else if (std::strcmp(value, "all") == 0 || std::strcmp(value, "on") == 0) {
        g_install_count = kProbeCount;
    } else {
        char* end = nullptr;
        const long n = std::strtol(value, &end, 10);
        if (end == value || n < 0) return true;  // ours, but unparseable: leave the default
        g_install_count = static_cast<std::size_t>(n) < kProbeCount
                              ? static_cast<std::size_t>(n)
                              : kProbeCount;
    }
    if (count_out) *count_out = g_install_count;
    return true;
}

void init_probe_entries(std::uintptr_t exe_base, void (*log_line)(const char* line)) {
    g_exe_base = exe_base;
    g_log = log_line;
}

void install_probe_entries() {
#if defined(_WIN32)
    logf("probe: installing %u of %u (probes= in shim.cfg)",
         static_cast<unsigned>(g_install_count), static_cast<unsigned>(kProbeCount));
    for (std::size_t i = 0; i < g_install_count; ++i) {
        void* target = reinterpret_cast<void*>(g_exe_base + kProbes[i].rva);
        MH_STATUS s1 = MH_CreateHook(target, kProbes[i].stub, kProbes[i].trampoline);
        MH_STATUS s2 = s1 == MH_OK ? MH_EnableHook(target) : s1;
        g_installed[i] = (s2 == MH_OK);
        logf("probe: %s rva=0x%08x -> va=%p create=%s enable=%s", kProbes[i].name, kProbes[i].rva,
             target, MH_StatusToString(s1), MH_StatusToString(s2));
        if (!g_installed[i])
            logf("COVERAGE: probe %s NOT INSTALLED -- its call count is meaningless, not zero",
                 kProbes[i].name);
    }
#else
    for (std::size_t i = 0; i < kProbeCount; ++i) g_installed[i] = false;
#endif
}

void probe_entries_reset() {
    for (std::size_t i = 0; i < kProbeCount; ++i) g_calls[i] = 0;
}

std::size_t probe_entries_count() { return kProbeCount; }

std::size_t probe_entries_snapshot(ProbeEntryRow* out, std::size_t max) {
    const std::size_t n = kProbeCount < max ? kProbeCount : max;
    for (std::size_t i = 0; i < n; ++i) {
        out[i].name = kProbes[i].name;
        out[i].rva = kProbes[i].rva;
        out[i].calls = g_calls[i];
        out[i].total = g_total[i];
        out[i].installed = g_installed[i];
    }
    return n;
}

}  // namespace shim::hooks