--[[
Source: https://www.52pojie.cn/thread-2077472-1-1.html
	Cheat Engine Register Highlight Script
Features:
1. Menu Creation: Adds a "Register Highlight" main menu to the right-click menu of the CE disassembly window, with submenus for all general-purpose registers.
2. Interactive Operation: Click on a register name to toggle its highlight state (checked/unchecked).
3. Highlight Effect: Selected registers are displayed with yellow background + black text in the disassembly code.
4. Real-time Update: View is repainted immediately after toggling, no restart or refresh needed.
5. Case Compatibility: Matches register names in different case formats.

]]

-- Get the disassembler view object from the memory window
local disasmView = getMemoryViewForm().DisassemblerView

-- Define 64-bit and 32-bit register pairing relationships
local registerPairs = {
    {"RAX", "EAX"}, {"RBX", "EBX"}, {"RCX", "ECX"}, {"RDX", "EDX"},
    {"RSI", "ESI"}, {"RDI", "EDI"}, {"RBP", "EBP"}, {"RSP", "ESP"},
    {"R8", ""}, {"R9", ""}, {"R10", ""}, {"R11", ""},
    {"R12", ""}, {"R13", ""}, {"R14", ""}, {"R15", ""}
}

-- Create a state table to record whether each register needs highlighting, initially all false
local highlightStates = {}
for _, pair in ipairs(registerPairs) do
    highlightStates[pair[1]] = false
end

-- Create the main menu item and add it to the disassembly view's popup menu
local mainItem = createMenuItem(disasmView.PopupMenu)
mainItem.Caption = "Register Highlight"
mainItem.ImageIndex = 38

disasmView.PopupMenu.Items.add(mainItem)

-- Store references to submenu items for later operations
local subMenuItems = {}
for _, pair in ipairs(registerPairs) do
    local reg64, reg32 = pair[1], pair[2]

    -- Create a submenu item for each register
    local subItem = createMenuItem(mainItem)

    -- Set menu display text (show "RAX/EAX" format if a 32-bit register exists)
    if reg32 and reg32 ~= "" then
        subItem.Caption = string.format("%s/%s", reg64, reg32)
    else
        subItem.Caption = reg64
    end

    subItem.Checked = false  -- Initially unchecked

    mainItem.add(subItem)    -- Add submenu to main menu

    subMenuItems[reg64] = subItem  -- Save reference
end

-- Clear all highlights function
local separatorItem = createMenuItem(mainItem)
separatorItem.Caption = '-'  -- Separator
mainItem.add(separatorItem)

-- Create "Clear All" menu item
local clearAllItem = createMenuItem(mainItem)
clearAllItem.Caption = "Clear All Highlights"
mainItem.add(clearAllItem)

-- Core highlight function: process text after disassembly, add color markers for registers
local function highlightDisasm(sender, address, LastDisassembleData, result, description)
    for _, pair in ipairs(registerPairs) do
        local reg64, reg32 = pair[1], pair[2]

        -- If this register is marked for highlighting
        if highlightStates[reg64] then
            -- Create a case-insensitive regex pattern (e.g., "RAX" becomes "[Rr][Aa][Xx]")
            local pattern64 = reg64:gsub("%a", function(c)
                return "[" .. c:upper() .. c:lower() .. "]"
            end)
            -- Replace in the parameter string with color-coded format (B=background, C=foreground, N=reset)
            LastDisassembleData.parameters = LastDisassembleData.parameters:gsub(
                pattern64, '{BFFFF00}{C000000}%0{N}'
            )

            -- Perform the same operation for the 32-bit register
            if reg32 and reg32 ~= "" then
                local pattern32 = reg32:gsub("%a", function(c)
                    return "[" .. c:upper() .. c:lower() .. "]"
                end)
                LastDisassembleData.parameters = LastDisassembleData.parameters:gsub(
                    pattern32, '{BFFFF00}{C000000}%0{N}'
                )
            end
        end
    end
    return result, description
end

-- Set click event handlers for each submenu item
for reg64, menuItem in pairs(subMenuItems) do
    menuItem.OnClick = function(sender)
        sender.Checked = not sender.Checked  -- Toggle checked state

        if disasmView then
            disasmView.Repaint()  -- Refresh disassembly view
        end

        -- Update register highlight state
        highlightStates[reg64] = not highlightStates[reg64]
    end
end

-- "Clear All" click event handler
clearAllItem.OnClick = function(sender)
    -- Iterate over all register submenu items, uncheck and reset highlight states
    for reg64, menuItem in pairs(subMenuItems) do
        menuItem.Checked = false
        highlightStates[reg64] = false
    end

    if disasmView then
        disasmView.Repaint()  -- Refresh disassembly view
    end
end

-- Register the highlight function to the disassembler's PostDisassemble event, so it is automatically called after each disassembly
getVisibleDisassembler().OnPostDisassemble = highlightDisasm