#include "shim/hooks/research.h"

#include <cstdarg>
#include <cstdio>
#include <cstring>
#include <stdexcept>
#include <string>
#include <vector>

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

#include "game/sim/research.h"
#include "game/sim/species.h"
#include "generated/sots_addresses.h"
#include "mars/rng/mt19937.h"

namespace shim::hooks {

using trace::Tv;
namespace tv = trace::tv;

namespace {

namespace A = sots::addr;

constexpr std::size_t kNodeSize = A::TechNode_size;  // 0x34
constexpr std::size_t kRngSize = A::RNG_size;        // 0x9cc
constexpr std::size_t kTreeHeadSize = A::TechTree_off_Nodes + 0xc;  // owner + the node vector
constexpr std::size_t kMaxNodes = 8192;              // loop guard for a garbage vector header
constexpr std::size_t kMaxAlloc = 1024;
constexpr int kMtWords = mars::rng::MT19937::N;

using CostFn = int(SHIM_THISCALL*)(void* tree, void* node);

struct Env {
    std::uintptr_t exe_base = 0;
    void (*log_line)(const char*) = nullptr;
    CostFn cost = nullptr;
};
Env g_env;

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

// ---- safe pointer chasing (same guard the M2 describers use) -------------------------------

bool readable(const void* p, std::size_t n) {
    if (!p) return false;
    if (n == 0) return true;
#if defined(_WIN32)
    const char* c = static_cast<const char*>(p);
    const char* const end = c + n;
    while (c < end) {
        MEMORY_BASIC_INFORMATION mbi;
        if (!VirtualQuery(c, &mbi, sizeof mbi)) return false;
        if (mbi.State != MEM_COMMIT) return false;
        if (mbi.Protect & (PAGE_NOACCESS | PAGE_GUARD)) return false;
        const DWORD ok = PAGE_READONLY | PAGE_READWRITE | PAGE_WRITECOPY | PAGE_EXECUTE_READ |
                         PAGE_EXECUTE_READWRITE | PAGE_EXECUTE_WRITECOPY;
        if (!(mbi.Protect & ok)) return false;
        c = static_cast<const char*>(mbi.BaseAddress) + mbi.RegionSize;
    }
    return true;
#else
    return true;
#endif
}

std::int32_t word_at(const void* obj, std::size_t off) {
    std::int32_t v = 0;
    std::memcpy(&v, static_cast<const char*>(obj) + off, sizeof v);
    return v;
}
void set_word(void* obj, std::size_t off, std::int32_t v) {
    std::memcpy(static_cast<char*>(obj) + off, &v, sizeof v);
}
void* ptr_at(const void* obj, std::size_t off) {
    void* v = nullptr;
    std::memcpy(&v, static_cast<const char*>(obj) + off, sizeof v);
    return v;
}
void set_ptr(void* obj, std::size_t off, void* v) {
    std::memcpy(static_cast<char*>(obj) + off, &v, sizeof v);
}

// The x87 control word as the hooked call finds it. The precision-control field decides
// whether the generator's multiply and the odds division round to 53 or to 24 significand
// bits, which is the one remaining float-parity unknown; recording it settles it from a
// single run instead of guessing.
std::uint32_t fpu_control_word() {
#if defined(__i386__) || defined(__x86_64__)
    unsigned short cw = 0;
    __asm__ __volatile__("fnstcw %0" : "=m"(cw));
    return cw;
#else
    return 0;
#endif
}

// ---- per-call state --------------------------------------------------------------------
//
// ProcessResearch is called once per player from the single-threaded turn pass and is never
// re-entrant, so the mapping regions() builds can be handed to rebind()/ours() in statics --
// the same concession M1 makes, and with the same caveat: do not reuse this for a hook that
// can nest.

struct CallState {
    bool compare = false;             // set by rebind, consumed (and cleared) by ours
    std::uintptr_t rng_base = 0;      // the LIVE generator address, for next-pointer math
    std::size_t rng_region = 0;
    std::size_t ob_region = 1;
    std::vector<int> node_region;     // node index -> region index, -1 when the slot is null
    std::vector<std::string> names;   // stable storage for Region::name
    std::vector<void*> scratch_nodes;
};
CallState g_call;

// The tree's node vector, or false when the header does not look like one.
bool tree_nodes(const void* tree, std::vector<void*>& out) {
    out.clear();
    if (!readable(tree, kTreeHeadSize)) return false;
    void** begin = static_cast<void**>(ptr_at(tree, A::TechTree_off_Nodes));
    void** end = static_cast<void**>(ptr_at(tree, A::TechTree_off_Nodes + 4));
    if (!begin && !end) return true;
    const std::uintptr_t b = reinterpret_cast<std::uintptr_t>(begin);
    const std::uintptr_t e = reinterpret_cast<std::uintptr_t>(end);
    if (!begin || e < b || (b & 3) || (e & 3)) return false;
    const std::size_t n = (e - b) / sizeof(void*);
    if (n > kMaxNodes) return false;
    if (n && !readable(begin, n * sizeof(void*))) return false;
    out.assign(begin, begin + n);
    return true;
}

// The {TechDef* target, int points} entries the budget built.
struct AllocEntry {
    void* target = nullptr;
    int points = 0;
};
bool alloc_entries(const void* alloc, std::vector<AllocEntry>& out) {
    out.clear();
    if (!readable(alloc, 0xc)) return false;
    const char* begin = static_cast<const char*>(ptr_at(alloc, 0));
    const char* end = static_cast<const char*>(ptr_at(alloc, 4));
    if (!begin && !end) return true;
    if (!begin || end < begin) return false;
    const std::size_t n = static_cast<std::size_t>(end - begin) / A::ResearchAlloc_stride;
    if (n > kMaxAlloc) return false;
    if (n && !readable(begin, n * A::ResearchAlloc_stride)) return false;
    for (std::size_t i = 0; i < n; ++i) {
        const char* e = begin + i * A::ResearchAlloc_stride;
        AllocEntry a;
        a.target = ptr_at(e, 0);
        a.points = word_at(e, 4);
        out.push_back(a);
    }
    return true;
}

int tech_id_of(void* def) {
    if (!readable(def, 4)) return -1;
    return word_at(def, A::TechDef_off_TechId);
}

// ---- describers -------------------------------------------------------------------------

Tv describe_rng(const void* p, std::size_t, unsigned inline_max) {
    Tv s = tv::struct_();
    s.add("vptr", tv::ptr(ptr_at(p, 0)));
    // The whole untempered block, hashed (it is far larger than inline_max): equality of the
    // hash plus equality of `left` is exactly "the two generators are in the same place in
    // the same stream".
    s.add("mt", tv::bytes(static_cast<const char*>(p) + A::RNG_off_State,
                          static_cast<std::size_t>(kMtWords) * 4, inline_max));
    const std::int32_t left = word_at(p, A::RNG_off_Left);
    s.add("left", tv::i32(left));
    // `next` is a heap address, so it is reported as its index into mt -- which is what it
    // means, and what survives being written by a reimplementation.
    const std::uintptr_t next = reinterpret_cast<std::uintptr_t>(ptr_at(p, A::RNG_off_Next));
    const std::uintptr_t base = g_call.rng_base + A::RNG_off_State;
    std::int64_t index = -1;
    if (g_call.rng_base && next >= base) index = static_cast<std::int64_t>((next - base) / 4);
    s.add("next_index", tv::i64(index));
    return s;
}

Tv describe_node(const void* p, std::size_t, unsigned) {
    Tv s = tv::struct_();
    void* def = ptr_at(p, A::TechNode_off_Def);
    s.add("def", tv::ptr(def));
    s.add("tech_id", tv::i32(tech_id_of(def)));
    // +4..+0xc is the children vector (three words) and +0x10 / +0x30 are not modelled;
    // emitted as opaque pointers so a change still shows without creating a false divergence.
    s.add("kids_begin", tv::ptr(ptr_at(p, 0x4)));
    s.add("kids_end", tv::ptr(ptr_at(p, 0x8)));
    s.add("kids_cap", tv::ptr(ptr_at(p, 0xc)));
    s.add("unk10", tv::ptr(ptr_at(p, 0x10)));
    s.add("state", tv::i32(word_at(p, A::TechNode_off_State)));
    s.add("cost_rp", tv::i32(word_at(p, A::TechNode_off_CostRP)));
    s.add("progress", tv::i32(word_at(p, A::TechNode_off_Progress)));
    s.add("turn_available", tv::i32(word_at(p, 0x20)));
    s.add("turn_researched", tv::i32(word_at(p, 0x24)));
    s.add("order", tv::i32(word_at(p, 0x28)));
    s.add("flag", tv::i32(word_at(p, A::TechNode_off_Flag)));
    s.add("unk30", tv::ptr(ptr_at(p, 0x30)));
    return s;
}

Tv describe_i32(const void* p, std::size_t, unsigned) {
    Tv s = tv::struct_();
    s.add("v", tv::i32(word_at(p, 0)));
    return s;
}

// ---- the generator seen by ours ----------------------------------------------------------

struct ShimRandom final : sots::sim::IRandom {
    mars::rng::MT19937 gen;
    unsigned draws = 0;
    float NextFloat() override {
        ++draws;
        return gen.next_float();
    }
    std::uint32_t NextIntInclusive(std::uint32_t n) override {
        ++draws;
        return gen.next_int_inclusive(n);
    }
    std::uint32_t NextUInt32() override {
        ++draws;
        return gen.next_u32();
    }
};

}  // namespace

// ---- descriptor ---------------------------------------------------------------------------

void TechTreeProcessResearchHook::describe_args(std::vector<Tv>& out, void* tree, void* rng,
                                                void* alloc, int* overbudget) {
    out.push_back(tv::ptr(tree).named("tree"));
    void* owner = readable(tree, kTreeHeadSize) ? ptr_at(tree, A::TechTree_off_Owner) : nullptr;
    out.push_back(tv::ptr(owner).named("owner"));
    const int species =
        readable(owner, A::ServerPlayer_off_Species + 4) ? word_at(owner, A::ServerPlayer_off_Species) : -1;
    out.push_back(tv::i32(species).named("species"));
    std::vector<void*> nodes;
    const bool nodes_ok = tree_nodes(tree, nodes);
    out.push_back(tv::u32(static_cast<std::uint32_t>(nodes.size())).named("node_count"));
    if (!nodes_ok) out.push_back(tv::boolean(true).named("nodes_invalid"));

    out.push_back(tv::ptr(rng).named("rng"));
    out.push_back(tv::i32(readable(rng, kRngSize) ? word_at(rng, A::RNG_off_Left) : -1).named("rng_left_in"));

    std::vector<AllocEntry> entries;
    const bool alloc_ok = alloc_entries(alloc, entries);
    std::vector<Tv> items;
    items.reserve(entries.size());
    for (const AllocEntry& e : entries) {
        Tv t = tv::struct_();
        t.add("tech_id", tv::i32(tech_id_of(e.target)));
        t.add("points", tv::i32(e.points));
        items.push_back(std::move(t));
    }
    out.push_back(tv::list(std::move(items)).named("alloc"));
    if (!alloc_ok) out.push_back(tv::boolean(true).named("alloc_invalid"));
    out.push_back(tv::i32(readable(overbudget, 4) ? word_at(overbudget, 0) : 0).named("overbudget_in"));
    // The x87 precision mode in force for this call (see fpu_control_word above).
    out.push_back(tv::u32(fpu_control_word()).named("fpu_cw"));
}

void TechTreeProcessResearchHook::regions(std::vector<trace::Region>& out, void* tree, void* rng,
                                          void* alloc, int* overbudget) {
    (void)alloc;
    g_call = CallState{};
    g_call.rng_base = reinterpret_cast<std::uintptr_t>(rng);

    if (!readable(rng, kRngSize)) throw std::runtime_error("rng object not readable");
    if (!readable(overbudget, 4)) throw std::runtime_error("overbudget not readable");

    trace::Region r;
    r.name = "rng";
    r.ptr = rng;
    r.size = kRngSize;
    r.describe = &describe_rng;
    g_call.rng_region = out.size();
    out.push_back(r);

    r = trace::Region{};
    r.name = "overbudget";
    r.ptr = overbudget;
    r.size = sizeof(int);
    r.describe = &describe_i32;
    g_call.ob_region = out.size();
    out.push_back(r);

    std::vector<void*> nodes;
    if (!tree_nodes(tree, nodes)) throw std::runtime_error("tech-tree node vector not readable");
    // Reserve once: Region::name holds a pointer into these strings for the whole call, so the
    // vector must never reallocate afterwards.
    g_call.names.resize(nodes.size());
    g_call.node_region.assign(nodes.size(), -1);
    for (std::size_t i = 0; i < nodes.size(); ++i) {
        char buf[32];
        std::snprintf(buf, sizeof buf, "node[%u]", static_cast<unsigned>(i));
        g_call.names[i] = buf;
        void* p = nodes[i];
        if (!readable(p, kNodeSize)) continue;  // a null slot in the tree: nothing to compare
        g_call.node_region[i] = static_cast<int>(out.size());
        trace::Region n;
        n.name = g_call.names[i].c_str();
        n.ptr = p;
        n.size = kNodeSize;
        n.describe = &describe_node;
        out.push_back(n);
    }
}

TechTreeProcessResearchHook::Args TechTreeProcessResearchHook::rebind(trace::Scratch& s, void* tree,
                                                                     void* rng, void* alloc,
                                                                     int* overbudget) {
    (void)rng;
    (void)overbudget;
    g_call.scratch_nodes.assign(g_call.node_region.size(), nullptr);
    for (std::size_t i = 0; i < g_call.node_region.size(); ++i) {
        if (g_call.node_region[i] >= 0)
            g_call.scratch_nodes[i] = s.ptr(static_cast<std::size_t>(g_call.node_region[i]));
    }
    g_call.compare = true;
    // The tree pointer is passed through unchanged: ours only reads it (owner, node count) and
    // hands it to the game's own read-only Cost. Every node it writes is a scratch copy.
    return Args(tree, s.ptr(g_call.rng_region), alloc,
                static_cast<int*>(s.ptr(g_call.ob_region)));
}

void TechTreeProcessResearchHook::ours(void* tree, void* rng, void* alloc, int* overbudget) {
    using namespace sots::sim;
    const bool compare = g_call.compare;
    g_call.compare = false;  // replace-mode calls must not inherit a stale mapping

    std::vector<void*> live;
    if (!tree_nodes(tree, live)) throw std::runtime_error("tech-tree node vector not readable");
    const std::vector<void*>& nodes = compare ? g_call.scratch_nodes : live;
    if (nodes.size() != live.size()) throw std::runtime_error("scratch node mapping is stale");

    void* owner = ptr_at(tree, A::TechTree_off_Owner);
    const int species_index =
        readable(owner, A::ServerPlayer_off_Species + 4) ? word_at(owner, A::ServerPlayer_off_Species) : -1;
    const Species owner_species = (species_index >= 0 && species_index < kSpeciesCount)
                                      ? static_cast<Species>(species_index)
                                      : static_cast<Species>(kSpeciesCount);

    // Model every slot. A slot the tree does not have stays Hidden so the decay pass skips it,
    // exactly as the original's null check does.
    std::vector<ResearchNode> model(nodes.size());
    for (std::size_t i = 0; i < nodes.size(); ++i) {
        void* p = nodes[i];
        if (!p) {
            model[i].state = TechState::Hidden;
            model[i].cost = 0;
            continue;
        }
        model[i].state = static_cast<TechState>(word_at(p, A::TechNode_off_State));
        model[i].progress = word_at(p, A::TechNode_off_Progress);
        model[i].flag = static_cast<TechFlag>(word_at(p, A::TechNode_off_Flag));
        model[i].cost = 0;
    }

    std::vector<AllocEntry> raw;
    if (!alloc_entries(alloc, raw)) throw std::runtime_error("research allocation not readable");
    std::vector<ResearchAllocEntry> entries;
    entries.reserve(raw.size());
    for (const AllocEntry& e : raw) entries.push_back({tech_id_of(e.target), e.points});

    // Effective cost, from the game's own read-only TechTree::Cost, for every node the
    // original would ask about: the allocation targets and every available node.
    if (!g_env.cost) throw std::runtime_error("TechTree::Cost address not initialised");
    auto fill_cost = [&](int i) {
        if (i < 0 || static_cast<std::size_t>(i) >= nodes.size() || !nodes[i]) return;
        model[i].cost = g_env.cost(tree, nodes[i]);
    };
    for (const ResearchAllocEntry& e : entries) fill_cost(e.nodeIndex);
    for (std::size_t i = 0; i < nodes.size(); ++i)
        if (model[i].state == TechState::Available) fill_cost(static_cast<int>(i));

    // Seed our generator from the pre-call state so both implementations read one stream.
    ShimRandom rand;
    {
        std::uint32_t mt[mars::rng::MT19937::N];
        std::memcpy(mt, static_cast<const char*>(rng) + A::RNG_off_State, sizeof mt);
        const std::int32_t left = word_at(rng, A::RNG_off_Left);
        if (left < 0 || left > kMtWords) throw std::runtime_error("generator `left` out of range");
        rand.gen.load_state(mt, left);
    }

    const ResearchTurnResult r = ProcessResearchTurn(model, entries, owner_species, rand);

    // Write back exactly the words the original function itself writes.
    for (std::size_t i = 0; i < nodes.size(); ++i) {
        void* p = nodes[i];
        if (!p) continue;
        set_word(p, A::TechNode_off_Progress, model[i].progress);
        set_word(p, A::TechNode_off_Flag, static_cast<std::int32_t>(model[i].flag));
        set_word(p, A::TechNode_off_State, static_cast<std::int32_t>(model[i].state));
    }
    set_word(overbudget, 0, word_at(overbudget, 0) + r.overbudget);

    // ... and the generator, in the object's own layout. `next` is rebuilt against the LIVE
    // base so the describer's index arithmetic is the same for both sides.
    {
        std::uint8_t blob[mars::rng::MT19937::kStateBytes];
        rand.gen.save_state(blob);
        std::memcpy(static_cast<char*>(rng) + A::RNG_off_State, blob,
                    static_cast<std::size_t>(kMtWords) * 4);
        set_word(rng, A::RNG_off_Left, rand.gen.left());
        const std::uintptr_t base =
            (compare ? g_call.rng_base : reinterpret_cast<std::uintptr_t>(rng)) + A::RNG_off_State;
        set_ptr(rng, A::RNG_off_Next,
                reinterpret_cast<void*>(base + static_cast<std::uintptr_t>(rand.gen.index()) * 4));
    }
}

void init_research(std::uintptr_t exe_base, void (*log_line)(const char* line)) {
    g_env.exe_base = exe_base;
    g_env.log_line = log_line;
    g_env.cost = reinterpret_cast<CostFn>(exe_base + A::TechTree_Cost);
    logf("research: ProcessResearch hook ready (Cost=%p, node=0x%x, rng=0x%x, fpu_cw=0x%04x)",
         reinterpret_cast<void*>(g_env.cost), static_cast<unsigned>(kNodeSize),
         static_cast<unsigned>(kRngSize), fpu_control_word());
}

}  // namespace shim::hooks