#include "shim/hooks/fleet_movement.h"

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

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

#include "generated/sots_addresses.h"
#include "mars/rng/mt19937.h"
#include "shim/hooks/movement_inputs.h"

namespace shim::hooks {

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

namespace {

namespace A = sots::addr;

constexpr std::size_t kRngSize = A::RNG_size;
constexpr int kMtWords = mars::rng::MT19937::N;
constexpr std::size_t kMaxPlayers = 64;

using ResolveWaypointFn = void*(SHIM_THISCALL*)(void* fleet);
using RelationFn = int(SHIM_THISCALL*)(void* player, void* other);

struct Env {
    std::uintptr_t exe_base = 0;
    void (*log_line)(const char*) = nullptr;
    ResolveWaypointFn resolve_waypoint = 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);
}

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); }

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
}

void* rng_of(void* server) {
    if (!readable(server, A::StrategyServer_off_RNGPtr + 4)) return nullptr;
    void* r = ptr_at(server, A::StrategyServer_off_RNGPtr);
    return readable(r, kRngSize) ? r : nullptr;
}

// The front waypoint of a fleet, or null when the vector is empty.
const void* front_waypoint(const void* fleet, std::int32_t* count) {
    *count = 0;
    if (!readable(fleet, A::StarFleet_off_Flags + 4)) return nullptr;
    const char* first = static_cast<const char*>(ptr_at(fleet, A::StarFleet_off_Waypoints));
    const char* last = static_cast<const char*>(ptr_at(fleet, A::StarFleet_off_Waypoints + 4));
    if (!first || !last || last <= first) return nullptr;
    const std::size_t n = static_cast<std::size_t>(last - first) / A::Waypoint_stride;
    if (n > 4096 || !readable(first, A::Waypoint_stride)) return nullptr;
    *count = static_cast<std::int32_t>(n);
    return first;
}

std::vector<void*> fleet_ships(const void* fleet) {
    std::vector<void*> ships;
    void** begin = static_cast<void**>(ptr_at(fleet, A::StarFleet_off_Ships));
    void** end = static_cast<void**>(ptr_at(fleet, A::StarFleet_off_Ships + 4));
    if (!begin || !end || end < begin) return ships;
    const std::size_t n = static_cast<std::size_t>(end - begin);
    if (n > kMaxTracedShips * 8 || !readable(begin, n * sizeof(void*))) return ships;
    ships.assign(begin, begin + n);
    return ships;
}

std::vector<void*> server_players(void* server) {
    std::vector<void*> out;
    if (!readable(server, A::StrategyServer_off_Players + 8)) return out;
    void** begin = static_cast<void**>(ptr_at(server, A::StrategyServer_off_Players));
    void** end = static_cast<void**>(ptr_at(server, A::StrategyServer_off_Players + 4));
    if (!begin || !end || end < begin) return out;
    const std::size_t n = static_cast<std::size_t>(end - begin);
    if (n > kMaxPlayers || !readable(begin, n * sizeof(void*))) return out;
    out.assign(begin, begin + n);
    return out;
}

std::vector<void*> server_fleets(void* server) {
    std::vector<void*> out;
    if (!readable(server, A::StrategyServer_off_Fleets + 8)) return out;
    void** begin = static_cast<void**>(ptr_at(server, A::StrategyServer_off_Fleets));
    void** end = static_cast<void**>(ptr_at(server, A::StrategyServer_off_Fleets + 4));
    if (!begin || !end || end < begin) return out;
    const std::size_t n = static_cast<std::size_t>(end - begin);
    if (n > kMaxTracedFleets || !readable(begin, n * sizeof(void*))) return out;
    out.assign(begin, begin + n);
    return out;
}

// ---- per-call state ------------------------------------------------------------------------

struct MoveState {
    bool compare = false;
    FleetStepSnapshot snap;
    double step = 0;
    std::uintptr_t rng_base = 0;
    std::vector<std::uint32_t> mt_before;
    std::int32_t left_before = 0;
    bool have_rng = false;
    void* s_pos = nullptr;
    void* s_prev = nullptr;
    void* s_rng = nullptr;
    std::vector<void*> s_ship;
    std::vector<std::string> names;  // stable storage for Region::name
};
MoveState g_move;

struct PfmState {
    bool compare = false;
    std::vector<FleetSummary> fleets;
    std::vector<void*> players;
    std::vector<void*> scratch;
    std::vector<std::string> names;
    std::int32_t playerCount = 0;
};
PfmState g_pfm;

Tv describe_i32(const void* p, std::size_t, unsigned) {
    Tv s = tv::struct_();
    s.add("v", tv::i32(peek<std::int32_t>(p, 0)));
    return s;
}
Tv describe_f32(const void* p, std::size_t, unsigned) {
    Tv s = tv::struct_();
    s.add("v", tv::f32(peek<float>(p, 0)));
    return s;
}
Tv describe_vec3(const void* p, std::size_t, unsigned) {
    Tv s = tv::struct_();
    s.add("x", tv::f32(peek<float>(p, 0)));
    s.add("y", tv::f32(peek<float>(p, 4)));
    s.add("z", tv::f32(peek<float>(p, 8)));
    return s;
}
Tv describe_rng(const void* p, std::size_t, unsigned inline_max) {
    Tv s = tv::struct_();
    s.add("mt", tv::bytes(static_cast<const char*>(p) + A::RNG_off_State,
                          static_cast<std::size_t>(kMtWords) * 4, inline_max));
    s.add("left", tv::i32(peek<std::int32_t>(p, A::RNG_off_Left)));
    return s;
}

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();
    }
};

void capture_move(void* self, void* fleet, float dt) {
    MoveState m;
    FleetStepSnapshot& s = m.snap;
    s.dt = dt;
    s.fpuControlWord = fpu_control_word();
    if (!readable(fleet, A::StarFleet_off_Flags + 4)) {
        g_move = std::move(m);
        return;
    }
    s.fleetId = peek<std::int32_t>(fleet, A::StarFleet_off_Id);
    s.speed = peek<float>(fleet, A::StarFleet_off_Speed);
    s.heldFlag = (peek<std::uint32_t>(fleet, A::StarFleet_off_Flags) & 0x100u) ? 1 : 0;
    for (int i = 0; i < 3; ++i)
        s.pos[i] = peek<float>(fleet, A::StarFleet_off_Pos + std::size_t(i) * 4);

    void* owner = ptr_at(fleet, A::StarFleet_off_PID);
    if (readable(owner, A::ServerPlayer_off_CstT + 4)) {
        s.ownerIndex = peek<std::int32_t>(owner, A::ServerPlayer_off_PlyrIdx);
        s.castEfficiency = peek<float>(owner, A::ServerPlayer_off_CstE);
        s.castThreshold = peek<float>(owner, A::ServerPlayer_off_CstT);
    }

    const void* wp = front_waypoint(fleet, &s.waypointCount);
    if (wp) s.waypointType = peek<std::int32_t>(wp, A::Waypoint_off_Type);

    if (g_env.resolve_waypoint) {
        void* dest = g_env.resolve_waypoint(fleet);
        if (readable(dest, 0x24)) {
            s.destKind = peek<std::int32_t>(dest, 0x14);
            for (int i = 0; i < 3; ++i) s.dest[i] = peek<float>(dest, 0x18 + std::size_t(i) * 4);
        }
    }

    const std::vector<void*> ships = fleet_ships(fleet);
    s.shipCount = static_cast<std::int32_t>(ships.size());
    for (std::size_t i = 0; i < ships.size() && i < kMaxTracedShips; ++i) {
        if (!readable(ships[i], A::StarShip_off_MaxRange + 4)) continue;
        s.ships[i].range = peek<float>(ships[i], A::StarShip_off_Range);
        const std::uint32_t lo = peek<std::uint32_t>(ships[i], A::StarShip_off_Flags);
        s.ships[i].exempt = (lo & 0x1000u) ? 1 : 0;
    }

    void* rng = rng_of(self);
    if (rng) {
        m.rng_base = reinterpret_cast<std::uintptr_t>(rng);
        m.left_before = peek<std::int32_t>(rng, A::RNG_off_Left);
        m.mt_before.resize(static_cast<std::size_t>(kMtWords));
        std::memcpy(m.mt_before.data(), static_cast<const char*>(rng) + A::RNG_off_State,
                    m.mt_before.size() * 4);
        m.have_rng = true;
        s.rngLeftIn = m.left_before;
    }

    // A node-line waypoint's step comes from the stutter profile, which needs the node graph;
    // this milestone does not walk it, so the node-line case is reported and left to the
    // original. Every other case is `speed x dt`.
    m.step = sots::sim::StraightStep(s.speed, dt);
    g_move = std::move(m);
}

void capture_pfm(void* self) {
    PfmState p;
    p.players = server_players(self);
    p.playerCount = static_cast<std::int32_t>(p.players.size());
    for (void* f : server_fleets(self)) {
        FleetSummary fs;
        if (!readable(f, A::StarFleet_off_Flags + 4)) continue;
        fs.fleetId = peek<std::int32_t>(f, A::StarFleet_off_Id);
        fs.gateTraffic = peek<std::int16_t>(f, A::StarFleet_off_GateTraffic);
        void* owner = ptr_at(f, A::StarFleet_off_PID);
        if (readable(owner, A::ServerPlayer_off_PlyrIdx + 4))
            fs.ownerIndex = peek<std::int32_t>(owner, A::ServerPlayer_off_PlyrIdx);
        std::int32_t count = 0;
        const void* wp = front_waypoint(f, &count);
        if (wp) fs.waypointType = peek<std::int32_t>(wp, A::Waypoint_off_Type);
        p.fleets.push_back(fs);
    }
    g_pfm = std::move(p);
}

}  // namespace

void init_fleet_movement(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.resolve_waypoint =
        reinterpret_cast<ResolveWaypointFn>(exe_base + A::StarFleet_ResolveWaypoint);
}

// ---- MoveFleet -------------------------------------------------------------------------------

void StrategyServerMoveFleetHook::describe_args(std::vector<Tv>& out, void* self, void* fleet,
                                                float dt) {
    out.push_back(tv::ptr(self).named("server"));
    out.push_back(tv::ptr(fleet).named("fleet"));
    out.push_back(tv::f32(dt).named("dt"));
    out.push_back(tv::i32(g_move.snap.fleetId).named("fleet_id"));
    out.push_back(tv::i32(g_move.snap.ownerIndex).named("owner"));
    out.push_back(tv::i32(g_move.snap.waypointType).named("wpt_type"));
    out.push_back(tv::i32(g_move.snap.waypointCount).named("wpt_count"));
    out.push_back(tv::i32(g_move.snap.destKind).named("dest_kind"));
    out.push_back(tv::f32(g_move.snap.speed).named("speed"));
    out.push_back(tv::i32(g_move.snap.rngLeftIn).named("rng_left_in"));
    out.push_back(tv::u32(g_move.snap.fpuControlWord).named("fpu_cw"));
}

Tv StrategyServerMoveFleetHook::describe_ret(bool r) { return tv::boolean(r); }

void StrategyServerMoveFleetHook::regions(std::vector<trace::Region>& out, void* self,
                                          void* fleet, float dt) {
    capture_move(self, fleet, dt);
    if (!readable(fleet, A::StarFleet_off_Flags + 4)) return;

    trace::Region pos;
    pos.name = "pos";
    pos.ptr = static_cast<char*>(fleet) + A::StarFleet_off_Pos;
    pos.size = 12;
    pos.describe = &describe_vec3;
    out.push_back(pos);

    trace::Region prev;
    prev.name = "prev_pos";
    prev.ptr = static_cast<char*>(fleet) + A::StarFleet_off_PrvPos;
    prev.size = 12;
    prev.describe = &describe_vec3;
    out.push_back(prev);

    // One region per ship's remaining range, named so a diff points at the ship rather than
    // at a byte offset. `names` is sized once so no push_back can dangle a Region::name.
    const std::vector<void*> ships = fleet_ships(fleet);
    const std::size_t n = ships.size() < kMaxTracedShips ? ships.size() : kMaxTracedShips;
    g_move.names.clear();
    g_move.names.reserve(n);
    for (std::size_t i = 0; i < n; ++i) {
        char buf[32];
        std::snprintf(buf, sizeof buf, "ship[%u].range", static_cast<unsigned>(i));
        g_move.names.emplace_back(buf);
    }
    for (std::size_t i = 0; i < n; ++i) {
        if (!readable(ships[i], A::StarShip_off_Range + 4)) continue;
        trace::Region r;
        r.name = g_move.names[i].c_str();
        r.ptr = static_cast<char*>(ships[i]) + A::StarShip_off_Range;
        r.size = 4;
        r.describe = &describe_f32;
        out.push_back(r);
    }

    if (g_move.have_rng) {
        trace::Region r;
        r.name = "rng";
        r.ptr = reinterpret_cast<void*>(g_move.rng_base);
        r.size = kRngSize;
        r.describe = &describe_rng;
        out.push_back(r);
    }
}

StrategyServerMoveFleetHook::Args StrategyServerMoveFleetHook::rebind(trace::Scratch& s,
                                                                     void* self, void* fleet,
                                                                     float dt) {
    g_move.compare = true;
    g_move.s_pos = s.count() > 0 && s.size(0) >= 12 ? s.ptr(0) : nullptr;
    g_move.s_prev = s.count() > 1 && s.size(1) >= 12 ? s.ptr(1) : nullptr;
    g_move.s_ship.clear();
    std::size_t i = 2;
    for (; i < s.count() && s.size(i) == 4; ++i) g_move.s_ship.push_back(s.ptr(i));
    g_move.s_rng = i < s.count() && s.size(i) >= kRngSize ? s.ptr(i) : nullptr;
    return Args(self, fleet, dt);
}

bool StrategyServerMoveFleetHook::ours(void* self, void* fleet, float dt) {
    using H = trace::Hook<StrategyServerMoveFleetHook>;
    const bool compare = g_move.compare;
    g_move.compare = false;
    if (H::mode == trace::Mode::Replace) {
        // Replace is not offered: our side reproduces the step arithmetic but none of the
        // arrival handling, so feeding it to the game would strand every arriving fleet.
        static bool warned = false;
        if (!warned) {
            warned = true;
            logf("movement: replace mode is not supported for MoveFleet; "
                 "falling back to the original");
        }
        return H::original ? H::original(self, fleet, dt) : false;
    }
    if (!compare) return false;

    ShimRandom rng;
    if (g_move.have_rng) rng.gen.load_state(g_move.mt_before.data(), g_move.left_before);
    const FleetStepResult r = StepFleet(g_move.snap, g_move.step, rng);

    if (g_move.s_pos) {
        for (int i = 0; i < 3; ++i) {
            const float v = static_cast<float>(r.pos[i]);
            std::memcpy(static_cast<char*>(g_move.s_pos) + i * 4, &v, 4);
        }
    }
    // The previous position is only written when the fleet actually moved, and it takes the
    // ENTRY position -- not the new one.
    if (g_move.s_prev && r.moved != 0 && !(r.pos[0] == g_move.snap.pos[0] &&
                                           r.pos[1] == g_move.snap.pos[1] &&
                                           r.pos[2] == g_move.snap.pos[2])) {
        for (int i = 0; i < 3; ++i) {
            const float v = static_cast<float>(g_move.snap.pos[i]);
            std::memcpy(static_cast<char*>(g_move.s_prev) + i * 4, &v, 4);
        }
    }
    for (std::size_t i = 0; i < g_move.s_ship.size() && i < r.shipRanges.size(); ++i) {
        const float v = static_cast<float>(r.shipRanges[i]);
        std::memcpy(g_move.s_ship[i], &v, 4);
    }
    if (g_move.s_rng && g_move.have_rng) {
        char* dst = static_cast<char*>(g_move.s_rng);
        std::memcpy(dst + A::RNG_off_State, rng.gen.state(),
                    static_cast<std::size_t>(kMtWords) * 4);
        const std::int32_t left = rng.gen.left();
        std::memcpy(dst + A::RNG_off_Left, &left, 4);
        void* next = reinterpret_cast<void*>(g_move.rng_base + A::RNG_off_State +
                                             static_cast<std::size_t>(kMtWords - left) * 4);
        std::memcpy(dst + A::RNG_off_Next, &next, sizeof next);
    }
    return r.arrived;
}

// ---- ProcessFleetMovement --------------------------------------------------------------------

void StrategyServerProcessFleetMovementHook::describe_args(std::vector<Tv>& out, void* self) {
    out.push_back(tv::ptr(self).named("server"));
    out.push_back(tv::i32(g_pfm.playerCount).named("players"));
    out.push_back(tv::i32(static_cast<std::int32_t>(g_pfm.fleets.size())).named("fleets"));
    std::vector<Tv> summary;
    for (const FleetSummary& f : g_pfm.fleets) {
        Tv one = tv::struct_();
        one.add("id", tv::i32(f.fleetId));
        one.add("owner", tv::i32(f.ownerIndex));
        one.add("wpt_type", tv::i32(f.waypointType));
        one.add("gate_traffic", tv::i16(f.gateTraffic));
        summary.push_back(std::move(one));
    }
    out.push_back(tv::list(std::move(summary)).named("fleet_state"));
}

void StrategyServerProcessFleetMovementHook::regions(std::vector<trace::Region>& out,
                                                     void* self) {
    // Captured BEFORE the original runs, so the argument record shows the pre-move fleet
    // state; `ours` re-captures afterwards, because that is when the original sums the
    // traffic.
    capture_pfm(self);
    g_pfm.names.clear();
    g_pfm.names.reserve(g_pfm.players.size());
    for (std::size_t i = 0; i < g_pfm.players.size(); ++i) {
        char buf[32];
        std::snprintf(buf, sizeof buf, "gate_traffic[%u]", static_cast<unsigned>(i));
        g_pfm.names.emplace_back(buf);
    }
    for (std::size_t i = 0; i < g_pfm.players.size(); ++i) {
        if (!readable(g_pfm.players[i], A::ServerPlayer_off_GateTraffic + 4)) continue;
        trace::Region r;
        r.name = g_pfm.names[i].c_str();
        r.ptr = static_cast<char*>(g_pfm.players[i]) + A::ServerPlayer_off_GateTraffic;
        r.size = 4;
        r.describe = &describe_i32;
        out.push_back(r);
    }
}

StrategyServerProcessFleetMovementHook::Args
StrategyServerProcessFleetMovementHook::rebind(trace::Scratch& s, void* self) {
    g_pfm.compare = true;
    g_pfm.scratch.clear();
    for (std::size_t i = 0; i < s.count(); ++i)
        if (s.size(i) == 4) g_pfm.scratch.push_back(s.ptr(i));
    return Args(self);
}

void StrategyServerProcessFleetMovementHook::ours(void* self) {
    using H = trace::Hook<StrategyServerProcessFleetMovementHook>;
    const bool compare = g_pfm.compare;
    g_pfm.compare = false;
    if (H::mode == trace::Mode::Replace) {
        static bool warned = false;
        if (!warned) {
            warned = true;
            logf("movement: replace mode is not supported for ProcessFleetMovement; "
                 "falling back to the original");
        }
        if (H::original) H::original(self);
        return;
    }
    if (!compare) return;

    // Re-read the fleets: the original computes the gate-traffic total at the very end of the
    // pass, from the state the fleets are in after they have moved.
    capture_pfm(self);
    const std::vector<int> totals = GateTraffic(g_pfm.fleets, g_pfm.playerCount);
    for (std::size_t i = 0; i < g_pfm.scratch.size(); ++i) {
        const std::int32_t v = i < totals.size() ? totals[i] : 0;
        std::memcpy(g_pfm.scratch[i], &v, 4);
    }
}

}  // namespace shim::hooks