// Blitzschlag as chain lightning (rebalancing request of 2026-10-15).
//
// The game's Blitzschlag already chains (measured, docs/chain-lightning.md):
// the caster's art code creates a bolt object (type 872) at the target and
// sets it up with a cEventFX 0x60 whose [14] is the number of hits left in the
// chain (3 at level 1). The bolt hits its target; then, with hits left, its
// update (0x5b4640..) looks for the next target (0x6005d0 at 0x5b469f: near
// the current target, not the caster, not the current target), draws the arc
// (segment objects) and creates the next bolt with [14] - 1. Without a next
// target the chain ends, and between two enemies it bounces back and forth.
//
// This mod keeps all of that and decides:
// - hits per chain N = 3 + level / 5: the first bolt's [14];
// - the next target: the nearest enemy within half Blitzschlag's cast range (see
//   below; the game's own jump search used about 316) that this
//   chain has not hit yet; else another enemy; else the current target again
//   (a lone enemy takes all N hits);
// - chains B = 1 + level / 10: B - 1 more chains start at the first target,
//   each with a bolt at its own first jump target (created while the first
//   bolt is set up, as copies of its event; no arc for that first hop).

#include <cstdint>
#include <cstring>
#include <map>
#include <memory>
#include <string>
#include <vector>

#include <windows.h>

#include "game/engine/arts.h"
#include "game/engine/fx.h"
#include "game/engine/objects.h"
#include "game/engine/orders.h"
#include "re/hook.h"
#include "re/log.h"
#include "re/patch.h"

namespace mods::chainlightning {

uint32_t __fastcall NextTarget(void* manager, void*, const int32_t* position, uint32_t range, uint32_t caster,
                               uint32_t current, uint8_t* bolt);

namespace {

constexpr uint32_t kBoltObjectType = 0x368;  // 872
constexpr uint32_t kJumpSearchCall = 0x005b469f;  // in the bolt update
constexpr uint32_t kFindJumpTarget = 0x006005d0;  // cObjectManager::FindJumpTarget
// Blitzschlag's cast range: the art code (cCreature art switch 0x5222c0,
// case 0x21) lays the ray from the caster towards the target in segments of
// 30, at most 20, each on passable ground (0x636b60); only if the target is
// reached is the bolt created on it. A jump may go as far as a cast.
// The segment count is bound by the creature's array of segment ids
// (+0x263, 20 entries), so the range grows by longer segments (request of
// 2026-10-17: range +50 %, 600 -> 900): the ray's step and the segment length
// are patched in the art code, where the game uses its shared 30.0 constant.
constexpr float kRaySegmentLengthOriginal = 30.0f;
constexpr float kRaySegmentLength = 45.0f;
constexpr int kRaySegments = 20;
double CastRange() {
    return static_cast<double>(kRaySegmentLength) * kRaySegments;
}
// A jump between enemies reaches half as far as the cast (request of
// 2026-10-17: 900 felt too far for jumps).
double JumpRange() {
    return CastRange() / 2;
}
constexpr uint32_t kConstant30 = 0x0088ef18;  // float 30.0, shared by 72 instructions
constexpr uint32_t kRayRemainingMinusSegment = 0x00526db6;  // fsub [30.0]
constexpr uint32_t kRayLastSegmentCheck = 0x00526e0b;  // fcomp [30.0]: remaining shorter than a segment?
constexpr uint32_t kRaySegmentLengthInt = 0x00526e23;  // mov eax,30: a full segment's length
constexpr uint32_t kRayStep = 0x00526f3d;  // push 30.0f: advance the ray point
constexpr size_t kObjectOwnerPlayer = 0x39;  // u8, copied from the caster like the game does
constexpr size_t kCreaturePosition = 0x18;  // 4 dwords, as effects are placed
constexpr size_t kEventHitsLeft = 14;
constexpr DWORD kChainLifetimeMs = 10000;

struct Chain {
    std::vector<uint32_t> hit;  // targets in order
    DWORD started = 0;
};

struct Bolt {
    std::shared_ptr<Chain> chain;
    uint32_t hitsLeft = 0;
};

// The art the hero cast last (cast orders), for its level; a cast order arms
// the next bolt setup as a chain start.
uint16_t g_heroArt = 0;
bool g_armed = false;
// Bolt objects of our chains, by object id; and the chain the next bolt set
// up at a given target belongs to (filled by NextTarget and branch spawning).
std::map<uint32_t, Bolt> g_bolts;
struct Pending {
    std::shared_ptr<Chain> chain;
    uint32_t target;
};
std::vector<Pending> g_pending;

void OnOrder(void* creature, const game::CreatureOrder& order, void*) {
    if (order.command == game::kOrderUseArt && creature == game::GetPlayerHero()) {
        g_heroArt = static_cast<uint16_t>(order.art);
        g_armed = true;
    }
}

// Hits per chain and chains at an art level.
uint32_t Hits(unsigned level) {
    return 3 + level / 5;
}
unsigned Chains(unsigned level) {
    return 1 + level / 10;
}

unsigned ArtLevel(void* creature, uint16_t art) {
    uint8_t* entry = game::FindCooldownEntry(static_cast<uint8_t*>(creature) + game::kCreatureArts, art);
    return entry ? entry[game::kCooldownLevel] + entry[game::kCooldownBonusLevels] : 0;
}

// The nearest living enemy of `caster` within the jump range of `from`, not
// `current`; one not in `avoid` if there is any.
uint32_t Nearest(void* caster, void* from, uint32_t current, const std::vector<uint32_t>& avoid) {
    return game::NearestHostile(caster, from, JumpRange(), current, avoid.data(), avoid.size());
}

void Forget() {
    const DWORD now = GetTickCount();
    for (auto it = g_bolts.begin(); it != g_bolts.end();) {
        it = now - it->second.chain->started > kChainLifetimeMs ? g_bolts.erase(it) : std::next(it);
    }
    g_pending.clear();
}

// A bolt object at `target`, set up with a copy of `event` (while its strings
// are alive) with `hitsLeft`.
void SpawnBolt(const uint8_t* original, const uint32_t* event, uint32_t target, uint32_t hitsLeft) {
    auto* creature = static_cast<uint8_t*>(game::GetCreature(target));
    if (!creature) {
        return;
    }
    auto* bolt = static_cast<uint8_t*>(
        game::GetObjectById(game::CreateObjectAt(kBoltObjectType, creature + kCreaturePosition)));
    if (!bolt) {
        return;
    }
    bolt[kObjectOwnerPlayer] = original[kObjectOwnerPlayer];
    uint32_t copy[game::kFxEventDwords];
    std::memcpy(copy, event, sizeof(copy));
    copy[3] = target;
    std::memcpy(&copy[4], creature + kCreaturePosition, 16);
    copy[kEventHitsLeft] = hitsLeft;
    game::SendEvent(bolt, copy);
}

void BeforeFxEvent(uint8_t* fx, uint32_t* event) {
    if (event[1] != game::kFxEventBoltHit || *reinterpret_cast<uint32_t*>(fx + game::kObjectType) != kBoltObjectType ||
        !game::IsPlayerHeroId(event[2])) {
        return;
    }
    const uint32_t id = *reinterpret_cast<uint32_t*>(fx + game::kCreatureId);
    const uint32_t target = event[3];
    if (g_bolts.count(id)) {
        return;  // set up again later: already ours
    }
    // A jump of one of our chains (NextTarget or a branch said where it goes).
    for (auto it = g_pending.begin(); it != g_pending.end(); ++it) {
        if (it->target == target) {
            g_bolts[id] = {it->chain, event[kEventHitsLeft]};
            it->chain->hit.push_back(target);
            g_pending.erase(it);
            return;
        }
    }
    if (!g_armed) {
        return;
    }
    // The first bolt of a cast.
    g_armed = false;
    void* caster = game::GetCreature(event[2]);
    void* first = game::GetCreature(target);
    if (!caster || !first) {
        return;
    }
    Forget();
    const unsigned level = ArtLevel(caster, g_heroArt);
    const uint32_t hits = Hits(level);
    const unsigned chains = Chains(level);
    event[kEventHitsLeft] = hits;
    auto chain = std::make_shared<Chain>();
    chain->hit.push_back(target);
    chain->started = GetTickCount();
    g_bolts[id] = {chain, hits};
    // More chains: each starts with a bolt at its own first jump target.
    std::vector<uint32_t> taken{target};
    for (unsigned c = 1; c < chains && hits > 1; ++c) {
        const uint32_t next = Nearest(caster, first, target, taken);
        const uint32_t branchTarget = next ? next : target;
        taken.push_back(branchTarget);
        auto branch = std::make_shared<Chain>();
        branch->hit.push_back(target);
        branch->started = chain->started;
        g_pending.push_back({branch, branchTarget});
        SpawnBolt(fx, event, branchTarget, hits - 1);
    }
    re::Log("chainlightning: art %x level %u: %u chain(s) of %u hits from %x", g_heroArt, level, chains, hits,
            target);
}

// Called instead of the jump search (ecx = the object manager, four stack
// arguments); the bolt is in ebp.
__declspec(naked) void NextTargetThunk() {
    __asm {
        push ebp                 // the bolt
        push dword ptr [esp+14h] // current target
        push dword ptr [esp+14h] // caster
        push dword ptr [esp+14h] // range
        push dword ptr [esp+14h] // position
        call NextTarget
        ret 10h
    }
}

void Startup() {
    const bool ok = re::RedirectCall(kJumpSearchCall, kFindJumpTarget, reinterpret_cast<const void*>(&NextTargetThunk));
    re::Log("chainlightning: jump targets%s", ok ? "" : " (patch FAILED)");
    static const float segmentLength = kRaySegmentLength;
    const auto bits = [](float value) {
        uint32_t result;
        std::memcpy(&result, &value, sizeof(result));
        return result;
    };
    const bool range = re::RedirectOperand(kRayRemainingMinusSegment, kConstant30, &segmentLength) &&
                       re::RedirectOperand(kRayLastSegmentCheck, kConstant30, &segmentLength) &&
                       re::PatchInstructionDword(kRaySegmentLengthInt, static_cast<uint32_t>(kRaySegmentLengthOriginal),
                                                 static_cast<uint32_t>(kRaySegmentLength)) &&
                       re::PatchPushImm32(kRayStep, bits(kRaySegmentLengthOriginal), bits(kRaySegmentLength));
    re::Log("chainlightning: cast range %.0f%s", CastRange(), range ? "" : " (patch FAILED)");
}
RE_STARTUP(ChainLightningStartup, Startup);

struct Registration {
    Registration() {
        game::OnCreatureOrder(&OnOrder);
        game::BeforeFxEvent(&BeforeFxEvent);
    }
} g_registration;

} // namespace

// The next target of `bolt`'s chain (the game's search for bolts that are
// not ours).
uint32_t __fastcall NextTarget(void* manager, void*, const int32_t* position, uint32_t range, uint32_t caster,
                               uint32_t current, uint8_t* bolt) {
    const uint32_t id = bolt ? *reinterpret_cast<uint32_t*>(bolt + game::kCreatureId) : 0;
    const auto it = g_bolts.find(id);
    if (it == g_bolts.end()) {
        using FindFn = uint32_t(__fastcall*)(void*, void*, const int32_t*, uint32_t, uint32_t, uint32_t);
        return reinterpret_cast<FindFn>(kFindJumpTarget)(manager, nullptr, position, range, caster, current);
    }
    void* casterCreature = game::GetCreature(caster);
    void* from = game::GetCreature(current);
    if (!casterCreature || !from || it->second.hitsLeft <= 1) {
        return 0;
    }
    const std::shared_ptr<Chain> chain = it->second.chain;
    uint32_t next = Nearest(casterCreature, from, current, chain->hit);
    if (!next) {
        next = current;  // nobody else: hit the same enemy again
    }
    g_pending.push_back({chain, next});
    return next;
}

} // namespace mods::chainlightning

// The tooltip: the game's "Anzahl: X" line (its hit count) becomes the mod's
// two numbers, hits per chain and chains, in both the current level's section
// and the next level's (skill book). cCreatureArts::ArtTooltip (0x589270,
// thiscall on creature +0x3a8, (art, text, flag), ret 0xc) builds the text
// into a wide string {begin, end, capacity}; sections start with the level
// line, values are "Name: value" lines separated by '\n' (measured for art
// 0x21 at level 5: "...|Aktuelle Stufe: 5|Regeneration: 3.5s|Schaden: ...|
// Anzahl: 7", with flag 1 also "Nächste erlernte Stufe: 6" and its values).
namespace mods::chainlightning {
namespace {

constexpr uint16_t kArtBlitzschlag = 0x21;  // "Kettenblitz" in the tooltip
constexpr const wchar_t* kCurrentLevelLine = L"Aktuelle Stufe:";
constexpr const wchar_t* kNextLevelLine = L"Nächste erlernte Stufe:";
constexpr const wchar_t* kCountLine = L"Anzahl:";

// The game's wide string append (0x592ef0, thiscall (first, last, unused),
// ret 0xc); grows the buffer with the game's allocator.
void AppendText(void* text, const wchar_t* first, const wchar_t* last) {
    using Fn = void*(__fastcall*)(void*, void*, const wchar_t*, const wchar_t*, uint32_t);
    reinterpret_cast<Fn>(0x00592ef0)(text, nullptr, first, last, 0);
}

void SetText(void* text, const std::wstring& value) {
    auto** string = static_cast<wchar_t**>(text);
    if (string[0]) {
        string[1] = string[0];
        *string[0] = L'\0';
    }
    AppendText(text, value.data(), value.data() + value.size());
}

std::wstring Rewrite(const std::wstring& text, unsigned currentLevel) {
    std::wstring out;
    unsigned level = currentLevel;
    size_t start = 0;
    while (start < text.size()) {
        size_t end = text.find(L'\n', start);
        end = end == std::wstring::npos ? text.size() : end + 1;
        const std::wstring line = text.substr(start, end - start);
        start = end;
        if (line.find(kCurrentLevelLine) != std::wstring::npos) {
            level = currentLevel;
        } else if (line.find(kNextLevelLine) != std::wstring::npos) {
            level = currentLevel + 1;
        }
        const size_t count = line.find(kCountLine);
        if (count == std::wstring::npos) {
            out += line;
            continue;
        }
        const std::wstring prefix = line.substr(0, count);  // colour codes
        const std::wstring newline = line.back() == L'\n' ? L"\n" : L"";
        out += prefix + L"Treffer je Kette: " + std::to_wstring(Hits(level)) + L"\n";
        out += prefix + L"Ketten: " + std::to_wstring(Chains(level)) + L"\n";
        out += prefix + L"Reichweite: " + std::to_wstring(static_cast<int>(CastRange())) + L"\n";
        out += prefix + L"Sprungweite: " + std::to_wstring(static_cast<int>(JumpRange())) + newline;
    }
    return out;
}

void __fastcall Arts_Tooltip(void* self, void* edx, uint32_t art, void* text, uint32_t flag);
RE_HOOK(0x00589270, Arts_Tooltip);
void __fastcall Arts_Tooltip(void* self, void* edx, uint32_t art, void* text, uint32_t flag) {
    Arts_Tooltip_original(self, edx, art, text, flag);
    if (!text || static_cast<uint16_t>(art) != kArtBlitzschlag) {
        return;
    }
    auto* const* string = static_cast<wchar_t* const*>(text);
    if (!string[0] || string[1] <= string[0]) {
        return;
    }
    void* creature = static_cast<uint8_t*>(self) - game::kCreatureArts;
    const std::wstring before(string[0], string[1]);
    const std::wstring after = Rewrite(before, ArtLevel(creature, kArtBlitzschlag));
    if (after != before) {
        SetText(text, after);
    }
}

} // namespace
} // namespace mods::chainlightning