CalendarExtras sorting hack

-- New sorting buttons for the invite lists
-- @file-revision@

local core = CalendarExtras

local module = core:NewModule("Sorters")

 --[[ -------- Locals ------- ]]--

local overrideSort, sortReverse
local listIsDirty

local roles -- This will point to the roles-table in the savedvars file after OnInit

-- Todo: Maybe add hybrid to the list... should've made these into icons in the first place
local roleIndexToName = {
    "Tank", -- [1]
    "Heal", -- [2]
    "DPS",  -- [3]
    "?"     -- [4]
}
local roleColors = {
    { r = 0.78, g = 0.61, b = 0.43 }, -- Tank
    { r = 0.1,  g = 1,    b = 0.1 },  -- Healer
    { r = 1,    g = 0.1,  b = 0.1 },  -- DPS
    NORMAL_FONT_COLOR                 -- No role
}
local knownRolesPerClass = {
    MAGE    = 3,
    HUNTER  = 3,
    ROGUE   = 3,
    WARLOCK = 3
}

local massSelectModule

 --[[ ----------------------- ]]--

local SortList

do
    local statusTable, roleTable, classTable, inviteIndexTable = {}, {}, {}, {}

    -- Need to move the status codes around so they sort in more sensible way
    --  1.Invited, 5.Out, 3.Declined, 8.Notsignedup, 9.Tentative, 6.Standby, 7.Signedup  2.Accepted, 4.Confirmed
    --  1          2      3           4              5            6          7           8           9
    local statusSwap = {
--      1, 2, 3, 4, 5, 6, 7, 8, 9 =>
        1, 8, 3, 9, 2, 6, 7, 4, 5
    }
    local function Sorter(name1, name2)
        if overrideSort == "role" then
            if roleTable[name1] == roleTable[name2] then
                if statusTable[name1] == statusTable[name2] then
                    return name1 < name2
                else
                    return statusTable[name1] > statusTable[name2]
                end
            else
                return roleTable[name1] < roleTable[name2]
            end
        elseif overrideSort == "mystatus" then -- Hax to sort status by status -> role -> class -> name, rather than status -> name
            if statusTable[name1] == statusTable[name2] then
                if roleTable[name1] == roleTable[name2] then
                    if classTable[name1] == classTable[name2] then
                        return name1 < name2
                    else
                        return classTable[name1] < classTable[name2]
                    end
                else
                    return roleTable[name1] < roleTable[name2]
                end
            else
                return statusTable[name1] > statusTable[name2]
            end
        end
    end

    function SortList()
        local t, i = {}, 1
        while true do
            local name, level, _, class, inviteStatus = CalendarEventGetInvite(i)
            if not name then break end
            t[i] = name
            statusTable[name] = statusSwap[inviteStatus] or inviteStatus
            roleTable[name] = module:GetRole(name, class)
            classTable[name] = class
            inviteIndexTable[name] = i
            i = i + 1
        end
        table.sort(t, Sorter)
        for k, v in ipairs(t) do
            module.reSortedList[k] = inviteIndexTable[v]
        end
        wipe(statusTable)
        wipe(roleTable)
        wipe(classTable)
        wipe(inviteIndexTable)
    end
end

local function NewSorterOnClick(updateFunc, self, button, down)
    core.callbacks:Fire("Pre_Sort", true)
    overrideSort = "role"
    if #module.reSortedList == 0 then
        SortList()
    end
    if sortReverse == nil then -- Default to forward sort
        sortReverse = false
    else
        sortReverse = not sortReverse
    end
    PlaySound("igMainMenuOptionCheckBoxOn");
    CalendarContextMenu_Hide(CalendarViewEventInviteContextMenu_Initialize)
    CalendarContextMenu_Hide(CalendarCreateEventInviteContextMenu_Initialize)
    local _,  _,  _,  _,  _,  _,  _, _,  _,  _,  _,  _,  _,  _,  _,  _,  _,  _,  _,  _,  _,  _,  _, inviteType, calendarType = CalendarGetEventInfo()
    updateFunc(inviteType, calendarType)
    core.callbacks:Fire("Post_Sort", true)
end

local function UpdateInviteListButtons(_, frame)
    local name, class, role, roleButton, c, _
    for _, button in ipairs(frame.buttons) do
        if not button.inviteIndex then break end
        name, _, _, class = CalendarEventGetInvite(button.inviteIndex)
        roleButton = _G[button:GetName().."Role"]
        role = module:GetRole(name, class)
        c = roleColors[role]
        roleButton:SetText(roleIndexToName[role])
        roleButton:SetTextColor(c.r, c.g, c.b)
    end
end

local function ClearCustomSort(_, custom)
    if custom then return end
    wipe(module.reSortedList)
    overrideSort = nil
    sortReverse = nil
    listIsDirty = nil
end

function module:SetRole(name, role)
    if not name or not role then return end
    roles[name] = role ~= 4 and role or nil
    if overrideSort then
        wipe(module.reSortedList)
        SortList()
    end
end

function module:GetRole(name, class)
    if name then
        if roles[name] then
            return roles[name]
        end
    end
    if class and knownRolesPerClass[class] then
        return knownRolesPerClass[class]
    end
    return 4
end

function module:OnInitialize()
    -- Set up our roles database
    if not core.db.roles then
        core.db.rolesDBver = 1 -- Think I screwed up when I didn't add hybirds, might have to update the tables at some point, so lets assign a version to it
        core.db.roles = {}
    end
    roles = core.db.roles

    module.reSortedList = {}

    if core.modules["Mass Select"] then
        massSelectModule = core.modules["Mass Select"]
    end

    if core.modules.Export then
        core.modules.Export:AddFormatToken("ROLE", "Role", function(name, class) return roleIndexToName[self:GetRole(name, class)] end)
    end

    -- Create the new sorting buttons
    -- TODO: The direction arrow is overlapped by the Role text for some reason
    local viewBtn = CreateFrame("Button", "CalendarViewEventInviteListRoleSortButton", CalendarViewEventInviteList, "CalendarEventInviteSortButtonTemplate")
    local createBtn = CreateFrame("Button", "CalendarCreateEventInviteListRoleSortButton", CalendarCreateEventInviteList, "CalendarEventInviteSortButtonTemplate")

    viewBtn:SetText("Role")
    viewBtn:SetPoint("BOTTOM", CalendarViewEventInviteList, "TOP", 3, 0)
    viewBtn.criterion = "role"

    createBtn:SetText("Role")
    createBtn:SetPoint("BOTTOM", CalendarCreateEventInviteList, "TOP", 3, 0)
    createBtn.criterion = "role"

    local width = viewBtn:GetTextWidth() + _G["CalendarViewEventInviteListRoleSortButtonDirection"]:GetWidth()
    viewBtn:SetWidth(width)
    createBtn:SetWidth(width)

    viewBtn:SetScript("OnClick", function(...) NewSorterOnClick(CalendarViewEventInviteList_Update, ...) end)
    createBtn:SetScript("OnClick", function(...) NewSorterOnClick(CalendarCreateEventInviteList_Update, ...) end)

    CalendarViewEventInviteList.sortButtons["role"] = viewBtn
    CalendarCreateEventInviteList.sortButtons["role"] = createBtn

    -- This is kinda stupid, would be better to modify the templates, but I'm not sure how to do that before the frames get created, if that's even possible
    local function Modify(button)
        local buttonName = button:GetName()
        local --[[name,]] class, status, role = --[[_G[buttonName.."Name"],]] _G[buttonName.."Class"], _G[buttonName.."Status"], button:CreateFontString(buttonName.."Role", "BORDER", "GameFontNormalSmall")

        class:SetWidth(70)
        class:ClearAllPoints()
        class:SetPoint("RIGHT", -110, 0)
        class:SetJustifyH("RIGHT")

        status:ClearAllPoints()
        status:SetPoint("RIGHT", -35, 0)
        status:SetJustifyH("RIGHT")

        role:SetWidth(30)
        role:SetJustifyH("RIGHT")
        role:SetPoint("RIGHT")
    end
    for _, button in ipairs(CalendarViewEventInviteList.scrollFrame.buttons) do
        Modify(button)
    end
    for _, button in ipairs(CalendarCreateEventInviteList.scrollFrame.buttons) do
        Modify(button)
    end

    self:RegisterEvent("CALENDAR_UPDATE_INVITE_LIST", function()
        if overrideSort then
            listIsDirty = true
        end
    end)

    -- The main hook to swap player indices around
    core.CalendarEventGetInviteHandler = self.CalendarEventGetInvite

    -- Clear everything before lists update after event is opened
    core.RegisterCallback(self, "OpenEvent", ClearCustomSort)
    -- Clear everything when list is resorted
    core.RegisterCallback(self, "Pre_Sort", ClearCustomSort)

    -- Resort before list is updated if needed
    core.RegisterCallback(self, "Pre_Update", function()
        if listIsDirty then
            listIsDirty = nil
            wipe(module.reSortedList)
            SortList()
        end
    end)

    -- Go through the buttons after list is updated by the game
    --core.RegisterCallback(module, "Pre_Post_Update", UpdateInviteListButtons)
    core.RegisterCallback(self, "Post_Update", UpdateInviteListButtons)

    -- Hook the sort button functions
    self:RawHook("CalendarEventInviteList_AnchorSortButtons", true)

    -- Env proxy for the CalendarEventInviteList_UpdateSortButtons
    self:RawHook("CalendarEventGetInviteSortCriterion", true)

    local function StatusSortButtonOverride(frame, ...)
        core.callbacks:Fire("Pre_Sort", true)
        overrideSort = "mystatus"
        if #module.reSortedList == 0 then
            SortList()
        end
        if sortReverse == nil then -- Default to forward sort
            sortReverse = false
        else
            sortReverse = not sortReverse
        end
        PlaySound("igMainMenuOptionCheckBoxOn");
        CalendarContextMenu_Hide(CalendarViewEventInviteContextMenu_Initialize)
        CalendarContextMenu_Hide(CalendarCreateEventInviteContextMenu_Initialize)
        local _,  _,  _,  _,  _,  _,  _, _,  _,  _,  _,  _,  _,  _,  _,  _,  _,  _,  _,  _,  _,  _,  _, inviteType, calendarType = CalendarGetEventInfo()
        if frame == CalendarViewEventInviteListStatusSortButton then
            CalendarViewEventInviteList_Update(inviteType, calendarType)
        elseif frame == CalendarCreateEventInviteListStatusSortButton then
            CalendarCreateEventInviteList_Update(inviteType, calendarType)
        end
        core.callbacks:Fire("Post_Sort", true)
    end
    self:RawHookScript(CalendarViewEventInviteListStatusSortButton, "OnClick", StatusSortButtonOverride)
    self:RawHookScript(CalendarCreateEventInviteListStatusSortButton, "OnClick", StatusSortButtonOverride)

    -- Register our context menu entry
    core:NewContextMenuEntry({
        text = "Set Role",
        subMenu = "CalendarRoleContextMenu",
    })
end

 --[[ -------- Hooks -------- ]]--

-- CalendarEventGetInvite() hook, gets the real player index from our reSortedList table
function module.CalendarEventGetInvite(orig, ...)
    if overrideSort then
        local index = ...
        if sortReverse then
            index = #module.reSortedList - index + 1
        end
        if index and module.reSortedList[index] then
            return orig(module.reSortedList[index])
        end
    end
    return orig(...)
end

function module:CalendarEventInviteList_AnchorSortButtons(inviteList)
    local scrollFrame = inviteList.scrollFrame
    if not scrollFrame or not scrollFrame.buttons or not scrollFrame.buttons[1] then return end
    local inviteButtonName = scrollFrame.buttons[1]:GetName()

    local nameSortButton = inviteList.sortButtons.name
    if inviteList.partyMode then
        nameSortButton:SetPoint("LEFT", _G[inviteButtonName.."Name"], "LEFT")
    else
        nameSortButton:SetPoint("LEFT", _G[inviteButtonName.."PartyIcon"], "LEFT")
    end

    local classSortButton = inviteList.sortButtons.class
    classSortButton:SetPoint("RIGHT", _G[inviteButtonName.."Class"], "RIGHT")

    local statusSortButton = inviteList.sortButtons.status
    statusSortButton:SetPoint("RIGHT", _G[inviteButtonName.."Status"], "RIGHT")

    local roleSortButton = inviteList.sortButtons.role
    roleSortButton:SetPoint("RIGHT", _G[inviteButtonName.."Role"], "RIGHT")
end

function module:CalendarEventGetInviteSortCriterion()
    if overrideSort then
        return overrideSort, sortReverse
    else
        return self.hooks.CalendarEventGetInviteSortCriterion()
    end
end

 --[[ ------ Context Menu ------ ]]--

do
    local roleContextMenu = CreateFrame("Frame", "CalendarRoleContextMenu", CalendarContextMenu, "UIMenuTemplate")
    roleContextMenu:Hide()
    roleContextMenu:SetFrameStrata("FULLSCREEN")
    roleContextMenu:EnableMouse(true)
    roleContextMenu:SetScript("OnShow", function(self)
        self:Initialize()
    end)
    roleContextMenu:SetScript("OnEvent", function(self, event)
        if event == "CALENDAR_UPDATE_EVENT" then
            if self:IsShown() then
                self:Initialize()
            end
        end
    end)
    roleContextMenu:SetBackdropBorderColor(TOOLTIP_DEFAULT_COLOR.r, TOOLTIP_DEFAULT_COLOR.g, TOOLTIP_DEFAULT_COLOR.b)
    roleContextMenu:SetBackdropColor(TOOLTIP_DEFAULT_BACKGROUND_COLOR.r, TOOLTIP_DEFAULT_BACKGROUND_COLOR.g, TOOLTIP_DEFAULT_BACKGROUND_COLOR.b)
    roleContextMenu:RegisterEvent("CALENDAR_UPDATE_EVENT")
    roleContextMenu.parentMenu = "CalendarContextMenu"

    function roleContextMenu:Initialize()
        UIMenu_Initialize(self)

        local name = CalendarEventGetInvite(CalendarContextMenu.inviteButton.inviteIndex)

        for k, v in ipairs(roleIndexToName) do
            if v == "?" then
                UIMenu_AddButton(self)
            end
            UIMenu_AddButton(
                self,                   -- menu
                v == "Heal" and "Healer" or v == "?" and "Clear" or v, -- text
                nil,                    -- shortcut
                self.SetRole,           -- func
                nil,                    -- nested
                {name = name, role = k} -- value
            )
        end

        return UIMenu_FinishInitializing(self)
    end

    function roleContextMenu:SetRole()
        -- If the mass-select module is there and loaded, set the roles using it instead
        if massSelectModule then
            massSelectModule:RunOnSelected(function(...) module:SetRole(...) end, true, self.value.role)
        else
            module:SetRole(self.value.name, self.value.role)
        end

        -- Todo: Don't we need to mark the list dirty now?
        core:UpdateOpenList()
        CalendarContextMenu_Hide()
    end
end