// 7 weapon sets and 7 combat-art slots per inventory.
//
// cInventoryEntry keeps 5 of each in packed arrays (docs/skill-slots.md). The
// two extra slots live in a shadow inventory per inventory: a zeroed
// cInventoryEntry-sized block whose slots 0-1 are slots 5-6. The slot methods
// are hooked; for a slot >= 5 they run the original method on the shadow with
// slot - 5, so all side effects (item references, icon resources, the
// engine's art selection event) are the game's own code. The methods that
// walk all slots (Clear, ReplaceItemId, FixInventory) also run on the shadow.
//
// The shadow carries the inventory's player index (+0x1160) and flags
// (+0x1162, bit 2 = second art slot variant), copied before every call.
// Network events sent from the shadow would name slot - 5: multiplayer with
// more than 5 slots is not supported.

#include "mods/itemstack.h"
#include "mods/skillslots.h"

#include <cstring>
#include <mutex>

#include <windows.h>

#include "game/engine/items.h"
#include "game/engine/objects.h"
#include "game/engine/orders.h"
#include "re/hook.h"

namespace mods::skillslots {
namespace {

constexpr uint32_t kInventoryBlock = 0x00ad23c8;  // cInventoryManager: 32 inventories
constexpr unsigned kInventories = 32;
constexpr size_t kPlayerIndex = 0x1160;
constexpr size_t kFlags = 0x1162;

std::mutex g_lock;
uint8_t* g_shadows[kInventories] = {};

int IndexOf(const void* inventory) {
    const auto* base = *reinterpret_cast<const uint8_t* const*>(kInventoryBlock);
    const auto* p = static_cast<const uint8_t*>(inventory);
    if (!base || p < base || p >= base + kInventories * kInventorySize ||
        (p - base) % kInventorySize != 0) {
        return -1;
    }
    return static_cast<int>((p - base) / kInventorySize);
}

} // namespace

uint8_t* Inventory(unsigned index) {
    auto* base = *reinterpret_cast<uint8_t**>(kInventoryBlock);
    return base && index < kInventories ? base + index * kInventorySize : nullptr;
}

int HeroIndex() {
    void* hero = game::GetPlayerHero();
    const int index = hero ? game::CreatureField(hero, game::kCreatureId) : -1;
    return index >= 0 && index < game::kChestInventory ? index : -1;
}

uint8_t* HeroInventory() {
    const int hero = HeroIndex();
    return hero >= 0 ? Inventory(static_cast<unsigned>(hero)) : nullptr;
}

uint8_t* ShadowFor(const void* inventory, bool create) {
    const int index = IndexOf(inventory);
    if (index < 0) {
        return nullptr;
    }
    std::lock_guard guard(g_lock);
    uint8_t*& shadow = g_shadows[index];
    if (!shadow && create) {
        shadow = static_cast<uint8_t*>(VirtualAlloc(nullptr, kInventorySize, MEM_COMMIT | MEM_RESERVE,
                                                    PAGE_READWRITE));
    }
    if (shadow) {
        const auto* p = static_cast<const uint8_t*>(inventory);
        std::memcpy(shadow + kPlayerIndex, p + kPlayerIndex, 2);
        std::memcpy(shadow + kFlags, p + kFlags, 2);
    }
    return shadow;
}

} // namespace mods::skillslots

namespace game {

using mods::skillslots::kGameSlots;
using mods::skillslots::kSlots;
using mods::skillslots::ShadowFor;

namespace {

// Slot argument of a method (u16 in the game).
unsigned Slot(uint32_t value) {
    return value & 0xffff;
}

// Runs a slot method: slots 0-4 on the inventory, 5-6 on its shadow (with
// slot - 5; created on demand if `create`), others not at all (result 0).
template <class Original, class... Args>
auto ForSlot(Original original, void* self, uint32_t slot, bool create, Args... args) {
    if (Slot(slot) < kGameSlots) {
        return original(self, nullptr, slot, args...);
    }
    uint8_t* shadow = Slot(slot) < kSlots ? ShadowFor(self, create) : nullptr;
    if (!shadow) {
        return decltype(original(self, nullptr, slot, args...))();
    }
    return original(shadow, nullptr, Slot(slot) - kGameSlots, args...);
}

// Runs a method on the inventory and, if it has one, on its shadow.
template <class Original, class... Args>
void WithShadow(Original original, void* self, Args... args) {
    original(self, nullptr, args...);
    if (uint8_t* shadow = ShadowFor(self, false)) {
        original(shadow, nullptr, args...);
    }
}

} // namespace


// --- weapon sets: (set, hand, ...) ---------------------------------------------

uint32_t __fastcall Inv_GetWeaponSetItem(void* self, void* edx, uint32_t set, uint32_t hand);
RE_HOOK(0x005d86c0, Inv_GetWeaponSetItem);
uint32_t __fastcall Inv_GetWeaponSetItem(void* self, void*, uint32_t set, uint32_t hand) {
    return ForSlot(Inv_GetWeaponSetItem_original, self, set, false, hand);
}

void __fastcall Inv_SetWeaponSetItem(void* self, void* edx, uint32_t set, uint32_t hand, int id, int notify);
RE_HOOK(0x005d86f0, Inv_SetWeaponSetItem);
void __fastcall Inv_SetWeaponSetItem(void* self, void*, uint32_t set, uint32_t hand, int id, int notify) {
    ForSlot(Inv_SetWeaponSetItem_original, self, set, true, hand, id, notify);
}

void __fastcall Inv_SetWeaponSet(void* self, void* edx, uint32_t set, int left, int right, int notify);
RE_HOOK(0x005d8930, Inv_SetWeaponSet);
void __fastcall Inv_SetWeaponSet(void* self, void*, uint32_t set, int left, int right, int notify) {
    ForSlot(Inv_SetWeaponSet_original, self, set, true, left, right, notify);
}

uint32_t __fastcall Inv_TakeWeaponSetItem(void* self, void* edx, uint32_t set, uint32_t hand, void* out, int flag);
RE_HOOK(0x005d8be0, Inv_TakeWeaponSetItem);
uint32_t __fastcall Inv_TakeWeaponSetItem(void* self, void*, uint32_t set, uint32_t hand, void* out, int flag) {
    return ForSlot(Inv_TakeWeaponSetItem_original, self, set, false, hand, out, flag);
}

uint32_t __fastcall Inv_GetWeaponSetItemPos(void* self, void* edx, uint32_t set, uint32_t hand, void* out);
RE_HOOK(0x005d8d00, Inv_GetWeaponSetItemPos);
uint32_t __fastcall Inv_GetWeaponSetItemPos(void* self, void*, uint32_t set, uint32_t hand, void* out) {
    return ForSlot(Inv_GetWeaponSetItemPos_original, self, set, false, hand, out);
}

// --- combat-art slots: (slot, ...) ---------------------------------------------

uint32_t __fastcall Inv_GetArtSlotType(void* self, void* edx, uint32_t slot);
RE_HOOK(0x005d93e0, Inv_GetArtSlotType);
uint32_t __fastcall Inv_GetArtSlotType(void* self, void*, uint32_t slot) {
    return ForSlot(Inv_GetArtSlotType_original, self, slot, false);
}

uint32_t __fastcall Inv_GetArtSlotRef(void* self, void* edx, uint32_t slot);
RE_HOOK(0x005d9420, Inv_GetArtSlotRef);
uint32_t __fastcall Inv_GetArtSlotRef(void* self, void*, uint32_t slot) {
    return ForSlot(Inv_GetArtSlotRef_original, self, slot, false);
}

uint32_t __fastcall Inv_GetArtSlotData(void* self, void* edx, uint32_t slot);
RE_HOOK(0x005d9450, Inv_GetArtSlotData);
uint32_t __fastcall Inv_GetArtSlotData(void* self, void*, uint32_t slot) {
    return ForSlot(Inv_GetArtSlotData_original, self, slot, false);
}

void __fastcall Inv_SetArtSlot(void* self, void* edx, uint32_t slot, uint32_t type, uint32_t data, void* refs);
RE_HOOK(0x005d9480, Inv_SetArtSlot);
void __fastcall Inv_SetArtSlot(void* self, void*, uint32_t slot, uint32_t type, uint32_t data, void* refs) {
    ForSlot(Inv_SetArtSlot_original, self, slot, true, type, data, refs);
}

void __fastcall Inv_ClearArtSlot(void* self, void* edx, uint32_t slot);
RE_HOOK(0x005d9680, Inv_ClearArtSlot);
void __fastcall Inv_ClearArtSlot(void* self, void*, uint32_t slot) {
    ForSlot(Inv_ClearArtSlot_original, self, slot, false);
}

// SelectSlot(slot, kind 0 weapon set / 1 art slot, apply, notify): stores the
// selection; for an art slot also tells the engine which art the right mouse
// button uses (read from the slot entry, hence run on the shadow for 5-6).
void __fastcall Inv_SelectSlot(void* self, void* edx, uint32_t slot, uint32_t kind, int apply, int notify);
RE_HOOK(0x005d9820, Inv_SelectSlot);
void __fastcall Inv_SelectSlot(void* self, void*, uint32_t slot, uint32_t kind, int apply, int notify) {
    if (Slot(kind) > 1) {
        return;
    }
    ForSlot(Inv_SelectSlot_original, self, slot, true, kind, apply, notify);
    if (Slot(slot) >= kGameSlots && Slot(slot) < kSlots) {
        auto* selected = reinterpret_cast<uint16_t*>(static_cast<uint8_t*>(self) + mods::skillslots::kSelected);
        selected[Slot(kind)] = static_cast<uint16_t>(Slot(slot));
    }
}

// --- methods over all slots ----------------------------------------------------

void __fastcall Inv_Clear(void* self, void* edx, int full, uint32_t arg);
RE_HOOK(0x005d7d30, Inv_Clear);
void __fastcall Inv_Clear(void* self, void*, int full, uint32_t arg) {
    WithShadow(Inv_Clear_original, self, full, arg);
}

void __fastcall Inv_ReplaceItemId(void* self, void* edx, int oldId, int newId);
RE_HOOK(0x005d80e0, Inv_ReplaceItemId);
void __fastcall Inv_ReplaceItemId(void* self, void*, int oldId, int newId) {
    WithShadow(Inv_ReplaceItemId_original, self, oldId, newId);
    mods::itemstack::OnItemIdReplaced(static_cast<uint32_t>(oldId), static_cast<uint32_t>(newId));
}

void __fastcall Inv_FixInventory(void* self, void* edx);
RE_HOOK(0x005de8a0, Inv_FixInventory);
void __fastcall Inv_FixInventory(void* self, void*) {
    WithShadow(Inv_FixInventory_original, self);
}

} // namespace game