// 7 weapon-set and 7 combat-art slots on the taskbar.
//
// cUI_Taskbar2 keeps its slot controls in two packed arrays of 5 (+0x164
// weapon sets, +0x178 art slots) and walks them in loops of 5. The arrays move
// to the end of the taskbar object, which is allocated 0x40 bytes larger, with
// room for 7 each; every instruction addressing them, the loops and the slot
// count clamp are patched (tables in game/addresses.h, found by reviewing all
// of 0x6df340-0x6e6000, docs/skill-slots.md). The layout code then creates 7
// controls per row at the original spacing: weapon sets grow leftwards from
// the centre, art slots rightwards.
//
// The old arrays keep a copy of the first 5 controls, so code that was missed
// still sees valid controls.
//
// Combat-art slots are not selected any more: slot 0 is always the right
// mouse button's art (the game's "selected" slot), the others are cast with
// their keys (src/mods/skillkeys.cpp). Selecting another slot is refused
// unless a skill key asks for it. Slot 0 is drawn with the selected look and
// set apart from slots 1-6 by a gap; every ready art uses its selected icon.
//
// Recharging arts show their remaining cooldown (seconds, one decimal) on the
// slot, drawn with the game's font after RenderArtSlots.
// Combos too, from their charge (game::GetComboCooldown).
//
// Also: 7 entries in the slot-count-by-level table (6 at level 45, 7 at 65),
// and cCreature::AssignArtToSlot accepts slots up to 6.

#include "mods/skillbar.h"

#include <atomic>
#include <cstring>
#include <cwchar>

#include "game/addresses.h"
#include "game/engine/arts.h"
#include "game/engine/objects.h"
#include "game/ui/text.h"
#include "mods/skillslots.h"
#include "re/hook.h"
#include "re/log.h"
#include "re/patch.h"

namespace mods::skillbar {

void WeaponSetClickThunk();

namespace {

using namespace game::addr;
using mods::skillslots::kGameSlots;
using mods::skillslots::kSlots;

constexpr uint32_t kWeaponSlots = kTaskbarSize;  // new arrays behind the old object
constexpr uint32_t kSkillSlots = kWeaponSlots + kSlots * 4;
constexpr uint32_t kNewTaskbarSize = kSkillSlots + kSlots * 4 + 0x8;
constexpr int8_t kOldDistance = static_cast<int8_t>(kTaskbarSkillSlots - kTaskbarWeaponSlots);  // 0x14
constexpr int8_t kNewDistance = static_cast<int8_t>(kSkillSlots - kWeaponSlots);  // 0x1c

// cUI_Taskbar2 fields.
constexpr size_t kInventoryIndex = 0x158;  // u16: the hero's inventory
constexpr size_t kSlotCount = 0x15c;  // u16: slots shown, from SetSlotCount
// UI control: x at +0x24 (relative to the taskbar).
constexpr size_t kControlX = 0x24;

// Gap between slot 0 (right mouse button) and slots 1-6: half a slot.
constexpr int kRightButtonSlotGap = 0x21;
constexpr uint32_t kRightStripStart = 0x23a;

// Weapon set clicks (WeaponSetClickThunk).
constexpr uint32_t kWeaponSetEmptyCursorJump = 0x006e0ca7;
constexpr uint32_t kWeaponSetDrag = 0x006e10d8;

// Cooldown labels on the art slots.
constexpr uint16_t kCooldownFont = 1;
constexpr uint32_t kCooldownColor = 0xffffffff;
constexpr uint32_t kCooldownAlign = game::kTextCenter;

// (level, slots): the game's table plus slots 6 and 7.
constexpr uint16_t kSlotLevels[kSlots] = {0, 2, 8, 16, 30, 45, 65};

std::atomic<uint8_t*> g_taskbar{nullptr};
// A skill key is selecting a slot other than 0 (window thread).
bool g_keySelects = false;

// Where the cooldown labels go: computed on the first frame after the layout
// (render thread), fixed to the slots from then on.
game::UiRect g_labelRects[kSlots] = {};
bool g_labelRectsValid = false;

void Startup() {
    int failed = 0;
    failed += !re::PatchPushImm32(kTaskbarAllocSizePush, kTaskbarSize, kNewTaskbarSize);
    failed += re::PatchAll(kTaskbarWeaponSlotUsers, kTaskbarWeaponSlots, kWeaponSlots);
    failed += re::PatchAll(kTaskbarSkillSlotUsers, kTaskbarSkillSlots, kSkillSlots);
    failed += re::PatchAll(kTaskbarSlotDistanceDown, static_cast<uint8_t>(-kNewDistance));
    failed += re::PatchAll(kTaskbarSlotDistanceUp, static_cast<uint8_t>(kNewDistance));
    failed += re::PatchAll(kTaskbarSlotLoopBounds, static_cast<uint8_t>(kSlots));
    failed += re::PatchAll(kAssignArtMaxSlot, static_cast<uint8_t>(kSlots - 1));
    failed += re::ReplaceAll(kTaskbarSkillSelectionLook);
    failed += re::ReplaceAll(kTaskbarSignedHitTests);
    failed += !re::PatchInstructionDword(kTaskbarRightStripStart, kRightStripStart,
                                         kRightStripStart + kRightButtonSlotGap);
    failed += !re::RedirectJcc(kWeaponSetEmptyCursorJump, kWeaponSetDrag,
                               reinterpret_cast<const void*>(&WeaponSetClickThunk));
    static_assert(kOldDistance == 0x14);
    re::Log("skillbar: %u weapon sets / art slots on the taskbar (%d patches failed)", kSlots, failed);
}
RE_STARTUP(SkillBarStartup, Startup);

using SelectFn = void(__fastcall*)(void* taskbar, void* edx, uint32_t slot, int notify);

void Select(uint32_t function, unsigned slot) {
    if (uint8_t* taskbar = g_taskbar.load()) {
        reinterpret_cast<SelectFn>(function)(taskbar, nullptr, slot, 1);
    }
}

uint8_t* HeroInventory() {
    const uint8_t* taskbar = g_taskbar.load();
    return taskbar ? skillslots::Inventory(*reinterpret_cast<const uint16_t*>(taskbar + kInventoryIndex))
                   : nullptr;
}

} // namespace

void SelectSkillSlot(unsigned slot) {
    g_keySelects = slot != 0;
    Select(0x006e17d0, slot);  // cUI_Taskbar2::SelectSkillSlot
    g_keySelects = false;
}

void SelectWeaponSet(unsigned slot) {
    Select(0x006e1440, slot);  // cUI_Taskbar2::SelectWeaponSet
}

unsigned SelectedSkillSlot() {
    const uint8_t* inventory = HeroInventory();
    return inventory ? reinterpret_cast<const uint16_t*>(inventory + skillslots::kSelected)[1] : 0;
}

bool SkillSlotHasArt(unsigned slot) {
    const uint8_t* taskbar = g_taskbar.load();
    uint8_t* inventory = HeroInventory();
    if (!taskbar || !inventory || slot >= kSlots ||
        slot >= *reinterpret_cast<const uint16_t*>(taskbar + kSlotCount) ||
        !reinterpret_cast<void* const*>(taskbar + kSkillSlots)[slot]) {
        return false;
    }
    using GetArtSlotRefFn = uint32_t(__fastcall*)(void* inventory, void* edx, uint32_t slot);
    return reinterpret_cast<GetArtSlotRefFn>(0x005d9420)(inventory, nullptr, slot) != 0;
}

bool KeySelects() {
    return g_keySelects;
}

void OnTaskbar(void* taskbar) {
    g_taskbar = static_cast<uint8_t*>(taskbar);
}

// After the layout: slots 1-6 move right, away from slot 0; the first 5
// controls of each row are copied into the old arrays.
void AfterLayout(void* taskbar) {
    auto* p = static_cast<uint8_t*>(taskbar);
    auto* const* skills = reinterpret_cast<uint8_t* const*>(p + kSkillSlots);
    for (unsigned slot = 1; slot < kSlots; ++slot) {
        if (skills[slot]) {
            *reinterpret_cast<int*>(skills[slot] + kControlX) += kRightButtonSlotGap;
        }
    }
    std::memcpy(p + kTaskbarWeaponSlots, p + kWeaponSlots, kGameSlots * 4);
    std::memcpy(p + kTaskbarSkillSlots, p + kSkillSlots, kGameSlots * 4);
    g_labelRectsValid = false;
}

void UpdateLabelRects(uint8_t* const* skills) {
    for (unsigned slot = 0; slot < kSlots; ++slot) {
        if (skills[slot]) {
            g_labelRects[slot] =
                game::CenterLineVertically(kCooldownFont, game::ControlScreenRect(skills[slot]));
        }
    }
    g_labelRectsValid = true;
}

// The remaining cooldown in seconds (at the hero's current regeneration
// rate, game::GetArtCooldown), to one decimal, centred on each art slot that
// is recharging (combat moves and spells: slot types 2 / 3; combos: 4,
// game::GetComboCooldown).
void RenderCooldowns(void* taskbar, void* renderer) {
    auto* p = static_cast<uint8_t*>(taskbar);
    uint8_t* inventory = skillslots::Inventory(*reinterpret_cast<const uint16_t*>(p + kInventoryIndex));
    void* hero = game::GetPlayerHero();
    if (!inventory || !hero) {
        return;
    }
    using SlotFn = uint32_t(__fastcall*)(void* inventory, void* edx, uint32_t slot);
    const auto getType = reinterpret_cast<SlotFn>(0x005d93e0);  // cInventoryEntry::GetArtSlotType
    const auto getRef = reinterpret_cast<SlotFn>(0x005d9420);  // cInventoryEntry::GetArtSlotRef
    auto* const* skills = reinterpret_cast<uint8_t* const*>(p + kSkillSlots);
    if (!g_labelRectsValid) {
        UpdateLabelRects(skills);
    }
    const unsigned shown = *reinterpret_cast<const uint16_t*>(p + kSlotCount);
    for (unsigned slot = 0; slot < kSlots && slot < shown; ++slot) {
        const uint32_t type = getType(inventory, nullptr, slot);
        const uint32_t ref = getRef(inventory, nullptr, slot);
        if (!skills[slot] || (type != 2 && type != 3 && type != game::kArtSlotTypeCombo) || !ref) {
            continue;
        }
        const game::ArtCooldown cooldown = type == game::kArtSlotTypeCombo
                                               ? game::GetComboCooldown(hero, ref)
                                               : game::GetArtCooldown(hero, game::ResolveArtId(hero, ref));
        if (cooldown.remaining <= game::kArtReadyCooldown) {
            continue;  // ready: no label, the game's icon shows it
        }
        wchar_t text[16];
        swprintf_s(text, L"%.1f", cooldown.seconds);
        game::DrawText(renderer, kCooldownFont, g_labelRects[slot], text, kCooldownColor, kCooldownAlign);
    }
}

// A click on a weapon set with an empty cursor (cUI_Taskbar2's mouse handling,
// `je 0x6e10d8` at 0x6e0ca7) takes the set's weapon onto the cursor, i.e.
// starts dragging it. Only a click on the selected set does that now; a click
// on another set just selects it.

bool __fastcall SelectInsteadOfDrag(void* taskbar, void*, const uint8_t* inventory, uint32_t set) {
    set &= 0xffff;
    const unsigned selected = reinterpret_cast<const uint16_t*>(inventory + skillslots::kSelected)[0];
    if (set == selected) {
        return false;
    }
    using SelectFn = void(__fastcall*)(void*, void*, uint32_t, int);
    reinterpret_cast<SelectFn>(0x006e1440)(taskbar, nullptr, set, 1);  // cUI_Taskbar2::SelectWeaponSet
    return true;
}

// Jumped to instead of 0x6e10d8 (taskbar in esi, inventory in ebp, the set
// at [esp+0B8h]).
__declspec(naked) void WeaponSetClickThunk() {
    __asm {
        mov eax, [esp+0B8h]
        push eax
        push ebp
        mov ecx, esi
        call SelectInsteadOfDrag
        test al, al
        jz drag
        mov al, 1
        mov edx, 0x006e124d  // handler exit, al = handled
        jmp edx
    drag:
        mov edx, 0x006e10d8  // the drag (kWeaponSetDrag)
        jmp edx
    }
}

uint16_t SlotCountForLevel(unsigned level) {
    uint16_t count = 1;
    for (unsigned i = 0; i < kSlots; ++i) {
        if (level >= kSlotLevels[i]) {
            count = static_cast<uint16_t>(i + 1);
        }
    }
    return count;
}

} // namespace mods::skillbar

namespace game {

// cUI_Taskbar2's layout (0x6e4ba0, from the constructor): creates the slot
// controls.
void __fastcall UiTaskbar_Layout(void* self, void* edx);
RE_HOOK(0x006e4ba0, UiTaskbar_Layout);
void __fastcall UiTaskbar_Layout(void* self, void* edx) {
    mods::skillbar::OnTaskbar(self);
    UiTaskbar_Layout_original(self, edx);
    mods::skillbar::AfterLayout(self);
}

// cUI_Taskbar2::SetSlotCount (0x6dffa0): shows the first n slots of each row.
void __fastcall UiTaskbar_SetSlotCount(void* self, void* edx, uint32_t count);
RE_HOOK(0x006dffa0, UiTaskbar_SetSlotCount);
void __fastcall UiTaskbar_SetSlotCount(void* self, void* edx, uint32_t count) {
    mods::skillbar::OnTaskbar(self);
    UiTaskbar_SetSlotCount_original(self, edx, count);
}

// cUI_Taskbar2::SelectSkillSlot (0x6e17d0): makes combat-art slot `slot` the
// right mouse button's art. Clicks on the bar, the mouse wheel, COMBINE_SLOTS
// and the slot count check all end up here; only a skill key may select a
// slot other than 0.
void __fastcall UiTaskbar_SelectSkillSlot(void* self, void* edx, uint32_t slot, int notify);
RE_HOOK(0x006e17d0, UiTaskbar_SelectSkillSlot);
void __fastcall UiTaskbar_SelectSkillSlot(void* self, void* edx, uint32_t slot, int notify) {
    if (!mods::skillbar::KeySelects()) {
        slot = 0;
    }
    UiTaskbar_SelectSkillSlot_original(self, edx, slot, notify);
    // It also moves the selected slot background (orange) to the new slot;
    // redo the backgrounds, which mark slot 0 (kTaskbarSkillSelectionLook).
    if (self) {
        using UpdateSlotHighlightFn = void(__fastcall*)(void*, void*, int force);
        reinterpret_cast<UpdateSlotHighlightFn>(0x006df9a0)(self, nullptr, 1);
    }
}

// cUI_Taskbar2::RenderArtSlots (0x6e31e0): then the cooldown labels.
void __fastcall UiTaskbar_RenderArtSlots(void* self, void* edx, void* renderer);
RE_HOOK(0x006e31e0, UiTaskbar_RenderArtSlots);
void __fastcall UiTaskbar_RenderArtSlots(void* self, void* edx, void* renderer) {
    UiTaskbar_RenderArtSlots_original(self, edx, renderer);
    mods::skillbar::RenderCooldowns(self, renderer);
}

uint32_t __stdcall SlotCountForLevel(uint32_t level);
RE_HOOK(addr::SlotCountForLevel, SlotCountForLevel);
uint32_t __stdcall SlotCountForLevel(uint32_t level) {
    return mods::skillbar::SlotCountForLevel(level & 0xffff);
}

} // namespace game