#include "shim/hooks/system_output.h"

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

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

#include "game/sim/colony.h"
#include "game/sim/numeric.h"
#include "generated/sots_addresses.h"

namespace shim::hooks {

using trace::Tv;
namespace tv = trace::tv;
namespace A = sots::addr;

namespace {

constexpr std::size_t kSystemGuardSize = 0x2d8;   // the whole ServerSystem object
constexpr std::size_t kPopGroupStride = 0x18;     // {?, int type @+4, int species @+8, int64 @+0x10}
constexpr std::size_t kMaxPopGroups = 4096;
constexpr int kSpeciesSlots = 7;

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

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
}

template <class T>
T peek(const void* base, std::size_t off) {
    T v{};
    std::memcpy(&v, static_cast<const char*>(base) + off, sizeof v);
    return v;
}
void* ptr_at(const void* base, std::size_t off) { return peek<void*>(base, off); }

// A GlobalConst is reached through a pointer slot in .data holding the address of the storage
// word the data-file loader fills.
template <class T>
T global_const_via_slot(std::uintptr_t slot_rva, T fallback) {
    if (!g_env.exe_base) return fallback;
    void** slot = reinterpret_cast<void**>(g_env.exe_base + slot_rva);
    if (!readable(slot, sizeof(void*))) return fallback;
    void* storage = *slot;
    if (!readable(storage, sizeof(T))) return fallback;
    T v{};
    std::memcpy(&v, storage, sizeof v);
    return v;
}

// ... and a few are reached by their storage address directly.
template <class T>
T global_const_at(std::uintptr_t rva, T fallback) {
    if (!g_env.exe_base) return fallback;
    const void* p = reinterpret_cast<const void*>(g_env.exe_base + rva);
    if (!readable(p, sizeof(T))) return fallback;
    T v{};
    std::memcpy(&v, p, sizeof v);
    return v;
}

std::int64_t population_of(const void* pop, int groupType, int species) {
    if (!readable(pop, 0xc)) return 0;
    const char* begin = static_cast<const char*>(ptr_at(pop, 0x4));
    const char* end = static_cast<const char*>(ptr_at(pop, 0x8));
    if (!begin || !end || end < begin) return 0;
    const std::size_t n = static_cast<std::size_t>(end - begin) / kPopGroupStride;
    if (n > kMaxPopGroups || !readable(begin, n * kPopGroupStride)) return 0;
    std::int64_t sum = 0;
    for (std::size_t i = 0; i < n; ++i) {
        const char* e = begin + i * kPopGroupStride;
        if (peek<std::int32_t>(e, 0x4) != groupType) continue;
        if (peek<std::int32_t>(e, 0x8) != species) continue;
        sum += peek<std::int64_t>(e, 0x10);
    }
    return sum;
}

// ---- the tuning values these two formulas read out of the running process ----------------
//
// Every one of them is logged with every record, so a divergence can always be attributed to
// a value rather than to a guess about it.
struct OutputTuning {
    float stationBonusImperialOutput = 0;
    std::int32_t moraleIncreaseOutput = 0;
    float moraleIncreaseOutputMod = 0;
    std::int32_t moraleDecreaseOutput = 0;
    float moraleDecreaseOutputMod = 0;
    float slavesOutputMod = 0;
    // The population-type table the executable builds in code -- read live so the run either
    // confirms the initialiser reading or refutes it.
    float popTypeOut[3] = {0, 0, 0};
    std::int32_t popTypeMaxPop[3] = {0, 0, 0};
};

OutputTuning read_tuning() {
    OutputTuning t;
    t.stationBonusImperialOutput =
        global_const_via_slot<float>(A::GlobalConst_slot_STATION_BONUS_IMPERIAL_OUTPUT, 0.0f);
    t.moraleIncreaseOutput =
        global_const_via_slot<std::int32_t>(A::GlobalConst_slot_MORALE_INCREASE_OUTPUT, 0);
    t.moraleIncreaseOutputMod =
        global_const_via_slot<float>(A::GlobalConst_slot_MORALE_INCREASE_OUTPUT_MOD, 0.0f);
    t.moraleDecreaseOutput =
        global_const_via_slot<std::int32_t>(A::GlobalConst_slot_MORALE_DECREASE_OUTPUT, 0);
    t.moraleDecreaseOutputMod =
        global_const_via_slot<float>(A::GlobalConst_slot_MORALE_DECREASE_OUTPUT_MOD, 0.0f);
    t.slavesOutputMod = global_const_at<float>(A::GlobalConst_storage_SLAVES_OUTPUT_MOD, 0.0f);
    for (int i = 0; i < 3; ++i) {
        const std::uintptr_t row = A::PopTypeTable_base + std::uintptr_t(i) * 0x30;
        t.popTypeOut[i] = global_const_at<float>(row + 0x10, 0.0f);
        t.popTypeMaxPop[i] = global_const_at<std::int32_t>(row + 0x8, 0);
    }
    return t;
}

sots::sim::TuningTable to_tuning(const OutputTuning& t) {
    sots::sim::TuningTable out;
    out.STATION_BONUS_IMPERIAL_OUTPUT = t.stationBonusImperialOutput;
    out.MORALE_INCREASE_OUTPUT = t.moraleIncreaseOutput;
    out.MORALE_INCREASE_OUTPUT_MOD = t.moraleIncreaseOutputMod;
    out.MORALE_DECREASE_OUTPUT = t.moraleDecreaseOutput;
    out.MORALE_DECREASE_OUTPUT_MOD = t.moraleDecreaseOutputMod;
    out.SLAVES_OUTPUT_MOD = t.slavesOutputMod;
    return out;
}

Tv tuning_tv(const OutputTuning& t) {
    Tv s = tv::struct_();
    s.add("STATION_BONUS_IMPERIAL_OUTPUT", tv::f32(t.stationBonusImperialOutput));
    s.add("MORALE_INCREASE_OUTPUT", tv::i32(t.moraleIncreaseOutput));
    s.add("MORALE_INCREASE_OUTPUT_MOD", tv::f32(t.moraleIncreaseOutputMod));
    s.add("MORALE_DECREASE_OUTPUT", tv::i32(t.moraleDecreaseOutput));
    s.add("MORALE_DECREASE_OUTPUT_MOD", tv::f32(t.moraleDecreaseOutputMod));
    s.add("SLAVES_OUTPUT_MOD", tv::f32(t.slavesOutputMod));
    s.add("poptype0_out", tv::f32(t.popTypeOut[0]));
    s.add("poptype1_out", tv::f32(t.popTypeOut[1]));
    s.add("poptype2_out", tv::f32(t.popTypeOut[2]));
    s.add("poptype0_maxpop", tv::i32(t.popTypeMaxPop[0]));
    s.add("poptype1_maxpop", tv::i32(t.popTypeMaxPop[1]));
    s.add("poptype2_maxpop", tv::i32(t.popTypeMaxPop[2]));
    return s;
}

// ---- per-call state ----------------------------------------------------------------------
//
// Neither function nests and the strategic pass is single-threaded, so the snapshot taken in
// describe_args (before the original runs) reaches ours() through a static. Do not copy this
// into a re-entrant hook.

struct GroupState {
    bool ok = false;
    bool owned = false;
    bool independent = false;
    int morale = 0;
    OutputTuning tuning;
};
GroupState g_group;

struct TotalState {
    bool ok = false;
    sots::sim::BaseOutputInputs in;
    sots::sim::OutputModifiers mods;
    OutputTuning tuning;
    // diagnostics that never enter ours()
    std::int32_t systemIndex = 0;
    std::int32_t ownerSpecies = -1;
    std::int64_t slaveGroupPop = 0;
    std::int32_t addictionSlots = 0;
    float speciesResourceOutput = 0;
    std::int32_t speciesBaseDemand = 0;
};
TotalState g_total;

const void* species_def(int species) {
    if (!g_env.exe_base || species < 0 || species > 6) return nullptr;
    const void* p = reinterpret_cast<const void*>(A::SpeciesDefTable_base + g_env.exe_base +
                                                  std::uintptr_t(species) * 0x184);
    return readable(p, 0x54) ? p : nullptr;
}

int effective_species(const void* sys, const void* owner) {
    const void* indi = ptr_at(sys, A::ServerSystem_off_Indi);
    if (indi && readable(indi, 8)) return peek<std::int32_t>(indi, 4);
    if (owner && readable(owner, A::ServerPlayer_off_Species + 4))
        return peek<std::int32_t>(owner, A::ServerPlayer_off_Species);
    return -1;
}

// The Morale object at +0x11c is {vptr, int[7]}, so the per-species word is at +0x120 + 4*sp.
int morale_of(const void* sys, int species) {
    if (species < 0 || species >= kSpeciesSlots) return 0;
    const std::size_t off = A::ServerSystem_off_Morale + 4 + std::size_t(species) * 4;
    if (!readable(sys, off + 4)) return 0;
    return peek<std::int32_t>(sys, off);
}

}  // namespace

void init_system_output(std::uintptr_t exe_base, void (*log_line)(const char* line)) {
    g_env.exe_base = exe_base;
    g_env.log_line = log_line;
}

// ---- GroupOutput --------------------------------------------------------------------------

void ServerSystemGroupOutputHook::describe_args(std::vector<Tv>& out, void* self,
                                                std::int32_t groupType, std::int32_t species,
                                                double count) {
    GroupState st;
    st.tuning = read_tuning();
    if (readable(self, kSystemGuardSize)) {
        const void* owner = ptr_at(self, A::ServerSystem_off_PID);
        st.owned = owner != nullptr;
        st.independent = ptr_at(self, A::ServerSystem_off_Indi) != nullptr;
        st.morale = morale_of(self, species);
        st.ok = true;
        out.push_back(tv::i32(peek<std::int32_t>(self, A::ServerSystem_off_Idx)).named("sysIdx"));
    } else {
        out.push_back(tv::null().named("sysIdx"));
    }
    g_group = st;

    out.push_back(tv::ptr(self).named("this"));
    out.push_back(tv::i32(groupType).named("groupType"));
    out.push_back(tv::i32(species).named("species"));
    out.push_back(tv::f64(count).named("count"));
    out.push_back(tv::i32(st.morale).named("morale"));
    out.push_back(tv::boolean(st.owned).named("owned"));
    out.push_back(tv::boolean(st.independent).named("independent"));
    out.push_back(tuning_tv(st.tuning).named("tuning"));
}

Tv ServerSystemGroupOutputHook::describe_ret(double r) { return tv::f64(r); }

void ServerSystemGroupOutputHook::regions(std::vector<trace::Region>& out, void* self,
                                          std::int32_t, std::int32_t, double) {
    if (!readable(self, kSystemGuardSize)) return;
    trace::Region g;
    g.name = "guard:system";
    g.ptr = self;
    g.size = kSystemGuardSize;
    g.kind = trace::Region::Kind::Guard;
    out.push_back(g);
}

ServerSystemGroupOutputHook::Args ServerSystemGroupOutputHook::rebind(trace::Scratch&, void* self,
                                                                     std::int32_t groupType,
                                                                     std::int32_t species,
                                                                     double count) {
    return Args{self, groupType, species, count};
}

double ServerSystemGroupOutputHook::ours(void*, std::int32_t groupType, std::int32_t species,
                                         double count) {
    (void)species;
    if (!g_group.ok) return 0.0;
    if (groupType < 0 || groupType > 2) return 0.0;
    sots::sim::GroupOutputInputs in;
    in.group = static_cast<sots::sim::PopGroup>(groupType);
    // The count arrives as a double that the original formed from an int64; recovering the
    // integer keeps our arithmetic on the same path as the model's.
    in.count = static_cast<std::int64_t>(count);
    in.morale = g_group.morale;
    in.stations = 0;  // declared input boundary -- see coverage()
    in.owned = g_group.owned;
    in.independent = g_group.independent;
    return sots::sim::GroupOutput(in, to_tuning(g_group.tuning));
}

// ---- ComputeTotalOutput -------------------------------------------------------------------

void ServerSystemComputeTotalOutputHook::describe_args(std::vector<Tv>& out, void* self,
                                                       double overHarvestRate) {
    TotalState st;
    st.tuning = read_tuning();
    if (readable(self, kSystemGuardSize)) {
        const void* owner = ptr_at(self, A::ServerSystem_off_PID);
        const bool independent = ptr_at(self, A::ServerSystem_off_Indi) != nullptr;
        const int sp = effective_species(self, owner);

        st.systemIndex = peek<std::int32_t>(self, A::ServerSystem_off_Idx);
        st.ownerSpecies = sp;

        std::int64_t resAvail = peek<std::int32_t>(self, A::ServerSystem_off_Res);
        bool stripMines = false;
        if (owner && readable(owner, A::ServerPlayer_off_SetupOutputMult + 4)) {
            stripMines = peek<std::uint8_t>(owner, A::ServerPlayer_off_AMine) != 0;
            st.mods.playerOutMod = peek<float>(owner, A::ServerPlayer_off_OutMod);
            st.mods.rebOutMod = peek<float>(owner, A::ServerPlayer_off_RebOutMod);
            st.mods.scOutMod = peek<float>(owner, A::ServerPlayer_off_ScOutMod);
            st.mods.techOutMod = peek<float>(owner, A::ServerPlayer_off_SetupOutputMult);
        }
        if (stripMines) {
            resAvail += peek<std::int32_t>(self, A::ServerSystem_off_MRes);
            resAvail += peek<std::int32_t>(self, A::ServerSystem_off_ARes2);
        }

        sots::sim::BaseOutputInputs& in = st.in;
        in.imperialPopulation = std::int64_t(peek<std::int32_t>(self, A::ServerSystem_off_pbon)) +
                                std::int64_t(peek<std::int32_t>(self, A::ServerSystem_off_Pop));
        // Group 0 is credited to the system's effective species only; every other species
        // contributes nothing, which is why one int pair is the whole imperial term.
        if (sp < 0) in.imperialPopulation = 0;
        const void* pop2 = static_cast<const char*>(self) + A::ServerSystem_off_Pop2;
        const void* pbon2 = static_cast<const char*>(self) + A::ServerSystem_off_pbon2;
        for (int i = 0; i < kSpeciesSlots; ++i) {
            st.slaveGroupPop += population_of(pop2, 2, i) + population_of(pbon2, 2, i);
        }
        if (sp >= 0 && sp < kSpeciesSlots) {
            in.civilianPopulation = population_of(pop2, 1, sp) + population_of(pbon2, 1, sp);
            in.civilianMorale = morale_of(self, sp);
        }
        in.slavePopulation = 0;  // declared input boundary -- see coverage()
        in.stations = 0;         // declared input boundary
        in.independent = independent;
        in.transitResources = peek<std::int32_t>(self, A::ServerSystem_off_TRes);
        in.resourcesAvailable = resAvail;
        in.infra = peek<float>(self, A::ServerSystem_off_Infra);
        in.infraBonus = peek<float>(self, A::ServerSystem_off_ibon);
        in.overHarvestRate = overHarvestRate;

        if (const void* def = species_def(sp)) {
            st.speciesBaseDemand = peek<std::int32_t>(def, 0x4c);
            st.speciesResourceOutput = peek<float>(def, 0x50);
        }
        in.speciesBaseDemand = st.speciesBaseDemand;
        in.speciesResourceOutput = st.speciesResourceOutput;

        st.mods.owned = owner != nullptr;
        st.mods.rebelling = peek<std::int32_t>(self, A::ServerSystem_off_rbfl) != 0;
        st.mods.systemOutMod = peek<float>(self, A::ServerSystem_off_OutMod);
        st.mods.addictionPhase3 = false;  // declared input boundary
        st.ok = true;
    }
    g_total = st;

    out.push_back(tv::ptr(self).named("this"));
    out.push_back(tv::f64(overHarvestRate).named("SRoh"));
    out.push_back(tv::i32(st.systemIndex).named("sysIdx"));
    out.push_back(tv::i32(st.ownerSpecies).named("species"));
    out.push_back(tv::i64(st.in.imperialPopulation).named("imperialPop"));
    out.push_back(tv::i64(st.in.civilianPopulation).named("civilianPop"));
    out.push_back(tv::i64(st.slaveGroupPop).named("slaveGroupPop"));
    out.push_back(tv::i32(st.in.civilianMorale).named("civilianMorale"));
    out.push_back(tv::i64(st.in.resourcesAvailable).named("resAvail"));
    out.push_back(tv::i64(st.in.transitResources).named("TRes"));
    out.push_back(tv::f32(st.in.infra).named("Infra"));
    out.push_back(tv::f32(st.in.infraBonus).named("ibon"));
    out.push_back(tv::i32(st.speciesBaseDemand).named("speciesBaseDemand"));
    out.push_back(tv::f32(st.speciesResourceOutput).named("speciesResourceOutput"));
    out.push_back(tv::f32(static_cast<float>(st.mods.playerOutMod)).named("OutMod"));
    out.push_back(tv::f32(static_cast<float>(st.mods.systemOutMod)).named("sysOutMod"));
    out.push_back(tv::f32(static_cast<float>(st.mods.techOutMod)).named("setupOutMod"));
    out.push_back(tv::f32(static_cast<float>(st.mods.rebOutMod)).named("RebOutMod"));
    out.push_back(tv::f32(static_cast<float>(st.mods.scOutMod)).named("ScOutMod"));
    out.push_back(tv::boolean(st.mods.rebelling).named("rebelling"));
    out.push_back(tv::boolean(st.in.independent).named("independent"));
    out.push_back(tuning_tv(st.tuning).named("tuning"));
}

Tv ServerSystemComputeTotalOutputHook::describe_ret(double r) { return tv::f64(r); }

void ServerSystemComputeTotalOutputHook::regions(std::vector<trace::Region>& out, void* self,
                                                 double) {
    if (!readable(self, kSystemGuardSize)) return;
    trace::Region g;
    g.name = "guard:system";
    g.ptr = self;
    g.size = kSystemGuardSize;
    g.kind = trace::Region::Kind::Guard;
    out.push_back(g);
}

ServerSystemComputeTotalOutputHook::Args ServerSystemComputeTotalOutputHook::rebind(
    trace::Scratch&, void* self, double overHarvestRate) {
    return Args{self, overHarvestRate};
}

double ServerSystemComputeTotalOutputHook::ours(void*, double) {
    if (!g_total.ok) return 0.0;
    const sots::sim::TuningTable t = to_tuning(g_total.tuning);
    sots::sim::OutputModifiers m = g_total.mods;
    m.baseOutput = sots::sim::SystemBaseOutput(g_total.in, t);
    return sots::sim::TotalSystemOutputRaw(m, t);
}

}  // namespace shim::hooks