-- Armorlock by hugeblank, March 2023
-- A simple and extensible armor swapping software.
-- - Configurable via the "armorlock" namespace:
--   - usage: `set armorlock.elytra_id sc_goodies:white_elytra`
--   - storage inventory target (storage) - Required
--   - manipulator target (manipulator) - Required
--   - player target(username) - Required
--   - elytra item target (elytra_id) - Optional, defaults to "minecraft:elytra"
--   - off-ground duration before elytra swap triggers (air_ticks) - Optional, defaults to 8
--   - held item slot to lock elytra swapping (held_slot) - Optional, defaults to 1
-- - Builtin handling for elytra swapping and gold boots:
--   - Elytra:
--     - Swaps ON if off-ground for a configurable amount of ticks (air_ticks).
--     - Swaps OFF on contact with the ground, UNLESS holding shift
--     - Will NOT swap if selecting the first inventory slot
--   - Boots:
--     - Swaps ON when in the nether
--     - Swaps OFF when not in the nether
-- - Easy to extend with the createslotHandler function:
--   - Easily creates a way of tracking and toggling armor swapping


-- Feel free to modify to your need, but don't expect/assume I'll help in any specific circumstance :thumbs_up:

assert(chatbox, "Chatbox not connected")

local config = {
    manipulator = settings.get("armorlock.manipulator"),
    storage = settings.get("armorlock.storage"),
    username = settings.get("armorlock.username"),
    air_ticks = tonumber(settings.get("armorlock.air_ticks")) or 8,
    held_slot = tonumber(settings.get("armorlock.held_slot")) or 1,
    elytra_id = settings.get("armorlock.elytra_id") or "minecraft:elytra",
}
-- Just a big old block of assertions, don't mind it.
assert(config.manipulator, "manipulator id missing, set with: `set armorlock.manipulator <peripheral id>`")
assert(config.storage, "storage id missing, set with: `set armorlock.storage <peripheral id>`")
assert(peripheral.isPresent(config.manipulator), "Could not find manipulator "..tostring(config.manipulator))
assert(peripheral.isPresent(config.storage), "Could not find storage "..tostring(config.storage))
assert(tonumber(config.air_ticks), "air ticks must be a number")
config.air_ticks = tonumber(config.air_ticks)
assert(config.air_ticks > 0, "air ticks must be greater than 0")
assert(tonumber(config.held_slot), "held slot must be a number")
config.held_slot = tonumber(config.held_slot)
assert(config.held_slot > 0 and config.held_slot < 10, "held slot must be between 1 and 9 inclusive")
if not config.elytra_id:find("elytra") then
    print("warn: armorlock.elytra_id does not contain elytra in the name, swapping may not function as intended")
end

local manipulator = peripheral.wrap(config.manipulator)
assert(manipulator.getMetaOwner, "Introspection Module & Entity sensor required for this program.")

local function swapItem(equipmentinv, condition, storageslot, equipmentslot)
    assert(storageslot%2 == 1, "Storage slot must be odd!")
    if condition then
        equipmentinv.pushItems(config.storage, equipmentslot, 1, storageslot)
        equipmentinv.pullItems(config.storage, storageslot+1, 1, equipmentslot)
    else
        equipmentinv.pushItems(config.storage, equipmentslot, 1, storageslot+1)
        equipmentinv.pullItems(config.storage, storageslot, 1, equipmentslot)
    end
end

-- Autoincrement value to control inventory spaces.
-- Make sure you're providing an inventory with enough spaces for the amount of slotHandlers you need!
-- floor(size/2) = max amount of slotHandlers.
local slotoffset = 1

-- Whether swapping is locked. Slot handlers should respect this value.
local lock = false

-- Handler functions
local handlers = {}

--- Creates a pair of functions to get and toggle the state of the armor in the players inventory
--- Passes these functions into the handler function when initially invoked
--- @param equipment function[] the equipment pseudo-peripheral provided by the manipulator
--- @param equipmentslot int the slot number of the equipment to be swapped. See the comment on the line that defines the equipment variable for the mapping.
--- @param name string the MC identifier of the item to be worn in the "active" state, where the condition is being met.
--- @param handler function The handler for this equipment slot
local function createSlotHandler(equipment, equipmentslot, name, handler)
    local storageslot = slotoffset
    slotoffset = slotoffset + 2
    local list = equipment.list()
    local condition = list[equipmentslot] and list[equipmentslot].name == name
    local getState, toggle = function() 
        return condition 
    end, function()
        swapItem(equipment, condition, storageslot, equipmentslot)
        condition = not condition
        return condition
    end
    handlers[#handlers+1] = {
        getState, 
        toggle,
        function()
            handler(getState, toggle)
        end,
    }
end

while true do
    local ok, err = pcall(function()
        -- Reset values
        slotoffset, handlers, lock = 1, {}, false
        local equipment = manipulator.getEquipment() -- {mainhand, offhand, boots, leggings, chest, helmet} <- slot ID to use for each armor slot
        local list = equipment.list()

        createSlotHandler(equipment, 5, config.elytra_id, function(isWearingElytra, swap) -- Elytra
            local air_ticks = 0
            while true do
                if lock then
                    os.pullEvent("armor_lock")
                end
                local meta = manipulator.getMetaOwner()
                if meta.isAirborne then
                    air_ticks = air_ticks+1
                else
                    air_ticks = 0
                end
                local wearingElytra = isWearingElytra()
                if (not wearingElytra and meta.heldItemSlot ~= config.held_slot and (air_ticks == config.air_ticks or meta.isSwimming)) or (wearingElytra and not (meta.isAirborne or meta.isSwimming or meta.isSneaking) and not meta.isElytraFlying) then
                    swap()
                end
            end
        end)

        -- Change 3 to 6, and "minecraft:gold_boots" to "minecraft:gold_helmet" if you want to swap the helmet instead
        createSlotHandler(equipment, 3, "minecraft:gold_boots", function(isWearingGold, swap)  -- Boots
            while true do
                if lock then
                    os.pullEvent("armor_lock")
                end
                local _, user, origin, destination = os.pullEvent("world_change")
                if user == config.username then
                    local wearingGold = isWearingGold()
                    if (destination == "minecraft:the_nether" and not wearingGold) or (origin == "minecraft:the_nether" and wearingGold) then
                        swap()
                    end
                end
            end
        end)
        -- Add your own swap handler methods here. Handle swap conditions by passing a function into the waitForAny method below.

        chatbox.tell(config.username, "&aactive", "armorlock", nil, "format")

        local threads = {} -- Break out the handler functions to be thrown into parallel
        for _, functions in ipairs(handlers) do
            threads[#threads+1] = functions[3]
        end
        parallel.waitForAny(
            function() -- Disconnect handler
                while true do
                    local _, user = os.pullEvent("leave")
                    if user == config.username then
                        return
                    end
                end
            end,
            function() -- Armor lock toggle handler
                while true do
                    local _, user, command, args, data = os.pullEvent("command")
                    if user == config.username and command == "armorlock" then
                        lock = not lock
                        os.queueEvent("armor_lock")
                        for _, functions in ipairs(handlers) do
                            if functions[1]() ~= lock then
                                functions[2]()
                            end
                        end
                    end
                end
            end,
            table.unpack(threads)
        )
    end) -- The above functions have a high potential of erroring on player death or disconnect
    if not ok then
        printError(err)
        -- If you want to get a notification in chat of every error, uncomment the following line:
        -- chatbox.tell(config.username, "Error:\n&c"..tostring(err), "armorlock", nil, "format")
    end
    -- Check the playerlist first to make sure the player is still online (probably dead)
    local list = chatbox.getPlayerList()
    for i, v in ipairs(list) do
        list[v] = i
    end
    if not list[config.username] then -- If they're not online, wait for them to come back.
        while true do
            local _, user = os.pullEvent("join")
            if user == config.username then
                break
            end
        end
    else -- If they're presumably dead, wait a second before running again.
        sleep(1)
    end
    handlers = {}
end