local _, ns = ...

------------------------------------------------------------
-- Raid Swap Logic: multi-criteria balancer for two-split raids.
--
-- Public API
--   ns.RaidSwap:Suggest(roster, opts) -> { suggestions, splitA, splitB }
--     roster: list of { name, class, group, role, index, fullName, online }
--             (shape produced by RaidUI.GetRaidRoster)
--     opts:   { maxGroup = 6|8, piAssignments = table }
--
--   Suggestions: list of { p1 = player, p2 = player, reason = "..." }
--   Each suggestion is a 1-for-1 SwapRaidSubgroup between p1 and p2.
--
-- Criteria, in priority order:
--   1. Player count       (±1)
--   2. Tank distribution  (1 per split preferred)
--   3. Healer count       (±1)
--   4. Lust coverage      (each side has a Mage/Shaman/Evoker/Hunter)
--   5. Brez coverage      (each side has a Druid/DK/Warlock/Paladin)
--   6. Healer spec stack  (same spec not clustered on one side)
--   7. Melee/ranged       (rough DPS melee/ranged parity)
--   8. Debuff classes     (Monk / DH not all on one side)
--   9. PI pairing         (priest + assigned target same split)
--
-- Algorithm style is ported from the NSRT SetupManager heuristic
-- (criteria priority list + class/role spread). We emit minimal swaps
-- rather than NSRT's full re-arrange so mid-raid churn is low.
------------------------------------------------------------

ns.RaidSwap = ns.RaidSwap or {}
local SWAP = ns.RaidSwap

------------------------------------------------------------
-- Helpers
------------------------------------------------------------
local MELEE_CLASSES = {
    WARRIOR = true, ROGUE = true, DEATHKNIGHT = true, DEMONHUNTER = true,
    PALADIN = true, MONK = true, DRUID = true, SHAMAN = true,
}
local RANGED_CLASSES = {
    MAGE = true, WARLOCK = true, PRIEST = true, HUNTER = true, EVOKER = true,
}

-- Side resolver. Two split modes, matching NSRT's two mappings:
--   "odds"   → split A = odd groups (1,3,5...); split B = even (2,4,6...)
--   "halves" → split A = first half of groups (1..N/2);
--              split B = second half (N/2+1..N)
-- Returns true if the given group belongs to split A under the mode.
local function IsSideA(group, mode, maxGroup)
    if mode == "halves" then
        return group <= math.floor(maxGroup / 2)
    end
    -- default: odds/evens
    return group % 2 == 1
end

-- Pull specID from the shared RaidInspectData cache (populated by
-- RaidInspect.lua alongside tier info). May return 0 if we haven't
-- inspected this player yet — callers must handle that.
local function LookupSpecID(name)
    if not ns.RaidInspectData then return 0 end
    for _, data in pairs(ns.RaidInspectData) do
        if data.name == name then return data.specID or 0 end
    end
    return 0
end

local function EnrichPlayer(p)
    local specID = LookupSpecID(p.name)
    p.specID  = specID
    p.pos     = ns.SpecMeta and ns.SpecMeta:GetPos(specID, p.class, p.role) or 1
    p.canLust = ns.SpecMeta and ns.SpecMeta:CanLust(specID, p.class) or false
    p.canBrez = ns.SpecMeta and ns.SpecMeta:CanBrez(specID, p.class) or false
    return p
end

------------------------------------------------------------
-- Split builder
------------------------------------------------------------
local function NewSplitBucket()
    return {
        healers = {}, tanks = {}, dps = {},
        melee = {}, ranged = {},
        lust = {}, brez = {},
        all = {},
    }
end

function SWAP:BuildSplit(roster, maxGroup, splitMode)
    maxGroup = maxGroup or 6
    splitMode = splitMode or "odds"
    local split = {
        A = NewSplitBucket(),
        B = NewSplitBucket(),
        allPriests = {},
        mode       = splitMode,
        maxGroup   = maxGroup,
    }

    for _, p in ipairs(roster) do
        if p.group and p.group <= maxGroup then
            EnrichPlayer(p)
            local bucket = IsSideA(p.group, splitMode, maxGroup) and split.A or split.B
            table.insert(bucket.all, p)
            if p.role == "HEALER" then
                table.insert(bucket.healers, p)
            elseif p.role == "TANK" then
                table.insert(bucket.tanks, p)
            else
                table.insert(bucket.dps, p)
                if RANGED_CLASSES[p.class] then
                    table.insert(bucket.ranged, p)
                elseif MELEE_CLASSES[p.class] then
                    table.insert(bucket.melee, p)
                end
            end
            if p.canLust then table.insert(bucket.lust, p) end
            if p.canBrez then table.insert(bucket.brez, p) end
            if p.class == "PRIEST" then
                table.insert(split.allPriests, p)
            end
        end
    end

    return split
end

------------------------------------------------------------
-- Imbalance scoring — used by MakeHelpers to guarantee that every
-- suggested swap *strictly improves* the split state (soft criteria)
-- or resolves a hard constraint (tank/lust/brez missing, healer spec
-- stacked, PI pairing). Without this check the criteria ping-pong:
-- "Balance size" proposes a swap that happens to worsen melee/ranged,
-- then "Melee spread" proposes the reverse, and the user clicks in
-- circles.
------------------------------------------------------------
local function PlayerContrib(p)
    local c = { all = 1, tanks = 0, healers = 0, dps = 0, melee = 0, ranged = 0, lust = 0, brez = 0 }
    if p.role == "TANK" then c.tanks = 1
    elseif p.role == "HEALER" then c.healers = 1
    else
        c.dps = 1
        if RANGED_CLASSES[p.class] then c.ranged = 1
        elseif MELEE_CLASSES[p.class] then c.melee = 1
        end
    end
    if p.canLust then c.lust = 1 end
    if p.canBrez then c.brez = 1 end
    return c
end

local function BucketCounts(bucket)
    return {
        all     = #bucket.all,     tanks  = #bucket.tanks,
        healers = #bucket.healers, dps    = #bucket.dps,
        melee   = #bucket.melee,   ranged = #bucket.ranged,
        lust    = #bucket.lust,    brez   = #bucket.brez,
    }
end

local function SimulateSwap(cA, cB, pA, pB)
    local c1, c2 = PlayerContrib(pA), PlayerContrib(pB)
    local newA, newB = {}, {}
    for k, v in pairs(cA) do
        newA[k] = v - c1[k] + c2[k]
        newB[k] = cB[k] - c2[k] + c1[k]
    end
    return newA, newB
end

-- Weighted sum of split imbalances. Lower = more balanced. Weights
-- reflect how painful each disparity is at a Mythic-style boss:
-- tank counts matter most, healer counts second, raw body count third,
-- then melee/ranged density. Missing essentials (tank/lust/brez on a
-- side that has players) add a large constant so the optimizer eagerly
-- fixes them.
local function ImbalanceScore(cA, cB)
    local s = 0
    s = s + math.abs(cA.all     - cB.all)     * 3
    s = s + math.abs(cA.tanks   - cB.tanks)   * 10
    s = s + math.abs(cA.healers - cB.healers) * 6
    s = s + math.abs(cA.melee   - cB.melee)   * 2
    s = s + math.abs(cA.ranged  - cB.ranged)  * 2
    if cA.tanks == 0 and cA.all > 0 then s = s + 25 end
    if cB.tanks == 0 and cB.all > 0 then s = s + 25 end
    if cA.lust  == 0 and cA.all > 0 then s = s + 15 end
    if cB.lust  == 0 and cB.all > 0 then s = s + 15 end
    if cA.brez  == 0 and cA.all > 0 then s = s + 8  end
    if cB.brez  == 0 and cB.all > 0 then s = s + 8  end
    return s
end

-- Reasons whose value isn't captured by the count-based score. These
-- bypass the strict-improvement check because it would otherwise
-- reject a beneficial swap (e.g. putting a priest next to her PI
-- target doesn't move any bucket counter).
local HARD_REASONS = {
    ["Tank split"]     = true,
    ["Lust split"]     = true,
    ["Brez split"]     = true,
    ["PI pairing"]     = true,
    ["Healer stacked"] = true,
    ["Mystic Touch"]   = true,
    ["Chaos Brand"]    = true,
}

------------------------------------------------------------
-- Criterion helpers (closures share `used` + suggestion list + counts)
------------------------------------------------------------
local function MakeHelpers(suggestions, split, maxGroup, splitMode)
    local used = {}
    local cA = BucketCounts(split.A)
    local cB = BucketCounts(split.B)

    -- Per-group populations (1..8). Updated on accepted moves so later
    -- criteria can find groups with open slots (for further moves) or
    -- choose partners from a group that still has someone available.
    local perGroup = {}
    for g = 1, 8 do perGroup[g] = 0 end
    for _, p in ipairs(split.A.all) do
        perGroup[p.group] = (perGroup[p.group] or 0) + 1
    end
    for _, p in ipairs(split.B.all) do
        perGroup[p.group] = (perGroup[p.group] or 0) + 1
    end

    local H = {}

    -- First candidate from `list` we haven't already picked. Lists are
    -- expected to be in stable order (e.g. splitA.dps) so repeated calls
    -- march through alternative partners.
    function H.FindAvailable(list)
        for _, p in ipairs(list) do
            if not used[p.name] then return p end
        end
        return nil
    end

    -- Try to find a partner that matches `predicate` first, falling
    -- back to any available DPS. Keeps quality swaps on role parity.
    function H.FindMatching(list, predicate)
        for _, p in ipairs(list) do
            if not used[p.name] and predicate(p) then return p end
        end
        return nil
    end

    function H.Add(p1, p2, reason)
        if not p1 or not p2 then return false end
        if used[p1.name] or used[p2.name] then return false end
        if p1.name == p2.name then return false end

        -- Normalize so p1 is always the A-side player for simulation.
        local p1A = IsSideA(p1.group, splitMode, maxGroup)
        local p2A = IsSideA(p2.group, splitMode, maxGroup)
        if p1A == p2A then return false end  -- same split already
        if not p1A then p1, p2 = p2, p1 end

        local newA, newB = SimulateSwap(cA, cB, p1, p2)

        -- Soft criteria must strictly improve the overall imbalance
        -- score; hard constraints bypass the check.
        if not HARD_REASONS[reason] then
            if ImbalanceScore(newA, newB) >= ImbalanceScore(cA, cB) then
                return false
            end
        end

        table.insert(suggestions, { op = "swap", p1 = p1, p2 = p2, reason = reason })
        used[p1.name] = true
        used[p2.name] = true
        cA, cB = newA, newB
        return true
    end

    -- Move player `p` to `targetGroup`. Two flavors:
    --   * intra-split (source and target same side): used by the
    --     compaction criterion; split counts unchanged, only per-group
    --     populations shift.
    --   * cross-split (different sides): used by the count criterion
    --     to actually rebalance split sizes (SwapRaidSubgroup can't
    --     change split counts — it just exchanges two players).
    function H.AddMove(p, targetGroup, reason)
        if not p or used[p.name] then return false end
        if p.group == targetGroup then return false end

        local fromA = IsSideA(p.group, splitMode, maxGroup)
        local toA   = IsSideA(targetGroup, splitMode, maxGroup)

        if fromA ~= toA then
            -- Cross-split: player's bucket contribution leaves the from
            -- side and joins the to side. cA/cB shift accordingly.
            local c = PlayerContrib(p)
            if fromA then
                for k, v in pairs(cA) do
                    cA[k] = v - c[k]
                    cB[k] = cB[k] + c[k]
                end
            else
                for k, v in pairs(cB) do
                    cB[k] = v - c[k]
                    cA[k] = cA[k] + c[k]
                end
            end
        end

        -- Always update per-group populations.
        perGroup[p.group]     = (perGroup[p.group] or 0) - 1
        perGroup[targetGroup] = (perGroup[targetGroup] or 0) + 1

        table.insert(suggestions, {
            op = "move", p1 = p, targetGroup = targetGroup, reason = reason,
        })
        used[p.name] = true
        return true
    end

    function H.IsUsed(name) return used[name] == true end

    -- Live post-simulation counts. Looping criteria call this after
    -- each successful H.Add so they can re-check their condition with
    -- the updated state ("am I still imbalanced?") instead of the
    -- stale initial snapshot.
    function H.Counts() return cA, cB end

    -- Per-group populations after all accepted swaps/moves so far.
    -- Used by Count to find target groups with open slots and by the
    -- compaction criterion to consolidate low-to-high.
    function H.PerGroup() return perGroup end

    return H
end

------------------------------------------------------------
-- Criterion: total player count
-- Loops until the split is within 1 of balanced or we run out of
-- valid partners. Re-reads counts from H.Counts() each iteration so
-- previous swaps feed into the "still imbalanced?" decision.
------------------------------------------------------------
local function Criterion_Count(split, H)
    for _ = 1, 6 do  -- safety cap; 6 is plenty for a 30-man raid
        local cA, cB = H.Counts()
        local diff = cA.all - cB.all
        if math.abs(diff) < 2 then return end

        local heavy, light, heavyCount
        if diff > 0 then
            heavy, light, heavyCount = split.A, split.B, cA
        else
            heavy, light, heavyCount = split.B, split.A, cB
        end

        -- Prefer moving a surplus healer if the heavy side also has
        -- too many healers; otherwise move a DPS. Keeps size + role
        -- parity aligned in one swap instead of two.
        local healDiff = (diff > 0) and (cA.healers - cB.healers)
                                    or  (cB.healers - cA.healers)
        local moveP
        if healDiff >= 1 and heavyCount.healers > 2 then
            moveP = H.FindAvailable(heavy.healers)
        end
        moveP = moveP or H.FindAvailable(heavy.dps)
        local swapWith = H.FindAvailable(light.dps)
        if not (moveP and swapWith) then return end
        if not H.Add(moveP, swapWith, "Balance size") then return end
    end
end

------------------------------------------------------------
-- Criterion: tanks — ideal 1 per split
------------------------------------------------------------
local function Criterion_Tanks(split, H)
    local aT, bT = #split.A.tanks, #split.B.tanks
    if aT >= 2 and bT == 0 then
        local moveP = split.A.tanks[2]
        local swapWith = H.FindAvailable(split.B.dps) or H.FindAvailable(split.B.all)
        H.Add(moveP, swapWith, "Tank split")
    elseif bT >= 2 and aT == 0 then
        local moveP = split.B.tanks[2]
        local swapWith = H.FindAvailable(split.A.dps) or H.FindAvailable(split.A.all)
        H.Add(moveP, swapWith, "Tank split")
    end
end

------------------------------------------------------------
-- Criterion: healer count
-- Loops with live counts: a healer diff of 4 is fixed in two swaps
-- within a single refresh instead of "click → still 2 off → click".
------------------------------------------------------------
local function Criterion_Healers(split, H)
    for _ = 1, 4 do
        local cA, cB = H.Counts()
        local diff = cA.healers - cB.healers
        if math.abs(diff) < 2 then return end

        local heavy, light
        if diff > 0 then heavy, light = split.A, split.B
        else             heavy, light = split.B, split.A end

        -- Prefer moving a duplicate-spec healer if one exists, so we
        -- rebalance count AND de-stack specs in one swap.
        local specCount = {}
        for _, h in ipairs(heavy.healers) do
            if not H.IsUsed(h.name) then
                local key = h.specID ~= 0 and h.specID or h.class
                specCount[key] = (specCount[key] or 0) + 1
            end
        end
        local moveH
        for _, h in ipairs(heavy.healers) do
            if not H.IsUsed(h.name) then
                local key = h.specID ~= 0 and h.specID or h.class
                if specCount[key] >= 2 then moveH = h; break end
            end
        end
        moveH = moveH or H.FindAvailable(heavy.healers)
        local swapWith = H.FindAvailable(light.dps)
        if not (moveH and swapWith) then return end
        if not H.Add(moveH, swapWith, "Healer balance") then return end
    end
end

------------------------------------------------------------
-- Criterion: lust coverage — each side wants ≥1 lust source
------------------------------------------------------------
local function Criterion_Lust(split, H)
    local aL, bL = #split.A.lust, #split.B.lust
    if aL == 0 and bL >= 2 then
        -- Move a lust from B to A. Prefer one that ISN'T healer/tank.
        local moveP = H.FindMatching(split.B.lust, function(p) return p.role == "DAMAGER" end)
                   or split.B.lust[#split.B.lust]
        local swapWith = H.FindAvailable(split.A.dps) or H.FindAvailable(split.A.all)
        H.Add(moveP, swapWith, "Lust split")
    elseif bL == 0 and aL >= 2 then
        local moveP = H.FindMatching(split.A.lust, function(p) return p.role == "DAMAGER" end)
                   or split.A.lust[#split.A.lust]
        local swapWith = H.FindAvailable(split.B.dps) or H.FindAvailable(split.B.all)
        H.Add(moveP, swapWith, "Lust split")
    end
end

------------------------------------------------------------
-- Criterion: brez coverage — each side wants ≥1 brez
-- NSRT caps brez per side at 2; if one side has 3+ and the other 0 we
-- still only swap one — the other brez stays on the heavy side.
------------------------------------------------------------
local function Criterion_Brez(split, H)
    local aR, bR = #split.A.brez, #split.B.brez
    if aR == 0 and bR >= 2 then
        local moveP = H.FindMatching(split.B.brez, function(p) return p.role == "DAMAGER" end)
                   or split.B.brez[#split.B.brez]
        local swapWith = H.FindAvailable(split.A.dps) or H.FindAvailable(split.A.all)
        H.Add(moveP, swapWith, "Brez split")
    elseif bR == 0 and aR >= 2 then
        local moveP = H.FindMatching(split.A.brez, function(p) return p.role == "DAMAGER" end)
                   or split.A.brez[#split.A.brez]
        local swapWith = H.FindAvailable(split.B.dps) or H.FindAvailable(split.B.all)
        H.Add(moveP, swapWith, "Brez split")
    end
end

------------------------------------------------------------
-- Criterion: healer spec stacking
-- Only flag when a split has 2+ of one healer spec AND the other split
-- has none. Uses actual specID when available (captured by the PI-
-- extended inspect pass); falls back to class as proxy when spec data
-- is missing (e.g. a priest hasn't been inspected yet so we can't tell
-- Disc from Holy).
------------------------------------------------------------
local function HealerKey(h)
    if h.specID and h.specID ~= 0 then return h.specID end
    return h.class
end

local function Criterion_HealerSpecs(split, H)
    if math.abs(#split.A.healers - #split.B.healers) >= 2 then return end

    local function buildSpecMap(list)
        local m = {}
        for _, h in ipairs(list) do
            local k = HealerKey(h)
            m[k] = m[k] or {}
            table.insert(m[k], h)
        end
        return m
    end

    local specA = buildSpecMap(split.A.healers)
    local specB = buildSpecMap(split.B.healers)

    for key, list in pairs(specA) do
        if #list >= 2 and not specB[key] and #split.A.healers > 2 then
            local moveP = list[2]
            local swapWith = H.FindAvailable(split.B.dps)
            H.Add(moveP, swapWith, "Healer stacked")
        end
    end
    for key, list in pairs(specB) do
        if #list >= 2 and not specA[key] and #split.B.healers > 2 then
            local moveP = list[2]
            local swapWith = H.FindAvailable(split.A.dps)
            H.Add(moveP, swapWith, "Healer stacked")
        end
    end
end

------------------------------------------------------------
-- Criterion: melee vs ranged parity
-- Flags when one split has 3+ more melee DPS than the other (ranged
-- stacking gets harsh on cleave mechanics; melee stacking gets harsh
-- on single-target mechanics). Only triggers when both splits have
-- DPS available to swap.
------------------------------------------------------------
local function Criterion_MeleeRanged(split, H)
    for _ = 1, 4 do
        local cA, cB = H.Counts()
        local diff = cA.melee - cB.melee
        if math.abs(diff) < 3 then return end

        local moveP, swapWith
        if diff > 0 then
            moveP = H.FindAvailable(split.A.melee)
            swapWith = H.FindAvailable(split.B.ranged) or H.FindAvailable(split.B.dps)
        else
            moveP = H.FindAvailable(split.B.melee)
            swapWith = H.FindAvailable(split.A.ranged) or H.FindAvailable(split.A.dps)
        end
        if not (moveP and swapWith) then return end
        if not H.Add(moveP, swapWith, "Melee spread") then return end
    end
end

------------------------------------------------------------
-- Criterion: debuff-bringer classes (Mystic Touch / Chaos Brand)
-- Spreading Monks + DHs keeps the multiplicative debuff on both
-- splits. Same shape as the existing check but uses the helper lib.
------------------------------------------------------------
local DEBUFF_CLASSES = {
    { file = "MONK",        short = "Mystic Touch" },
    { file = "DEMONHUNTER", short = "Chaos Brand" },
}

local function Criterion_DebuffClasses(split, H)
    for _, dc in ipairs(DEBUFF_CLASSES) do
        local pA, pB = {}, {}
        for _, p in ipairs(split.A.all) do if p.class == dc.file then table.insert(pA, p) end end
        for _, p in ipairs(split.B.all) do if p.class == dc.file then table.insert(pB, p) end end
        if #pA >= 2 and #pB == 0 then
            local moveP = pA[2]
            local swapWith = H.FindAvailable(split.B.dps)
            H.Add(moveP, swapWith, dc.short)
        elseif #pB >= 2 and #pA == 0 then
            local moveP = pB[2]
            local swapWith = H.FindAvailable(split.A.dps)
            H.Add(moveP, swapWith, dc.short)
        end
    end
end

------------------------------------------------------------
-- Criterion: PI pairing (priest + target same split)
------------------------------------------------------------
local function Criterion_PIPairs(split, H, piAssignments, roster)
    if not piAssignments then return end
    local byName = {}
    for _, p in ipairs(roster) do byName[p.name] = p end

    for _, priest in ipairs(split.allPriests) do
        local targetName = piAssignments[priest.name]
        if targetName then
            local target = byName[targetName]
            if target then
                local priestSide = IsSideA(priest.group, split.mode, split.maxGroup)
                local targetSide = IsSideA(target.group, split.mode, split.maxGroup)
                if priestSide ~= targetSide then
                    H.Add(priest, target, "PI pairing")
                end
            end
        end
    end
end

------------------------------------------------------------
-- Criterion: raid-wide group compaction
-- Fill the lowest-numbered groups to 5, then spread the remainder
-- across the last two groups as evenly as possible. Applies to the
-- whole raid, not per-split — the goal is what the player sees in
-- the Blizzard raid display, which is a single ordered 1..maxGroup
-- list. Split-based balance (tanks, healers, etc.) is handled by
-- the swap criteria above and is not affected by this pass since
-- the prior criteria emit swaps, which leave per-group counts
-- unchanged.
--
-- k = number of groups filled to 5 (lowest-numbered). We always
-- leave at least 2 tail groups so a stray single player lands as
-- (1, 1) across two groups instead of (1, 0) — which is what drove
-- the 5+1 → 3+3 request.
--   k = min(maxGroup - 2, floor(N / 5))
-- with k = maxGroup when N = 5 * maxGroup exactly (no tail needed).
--
-- The tail gets ceil/floor(remainder / tailCount), the lowest-
-- numbered tail group taking the high slot when it doesn't divide
-- cleanly. Examples (maxGroup = 6):
--   N=22 → 5,5,5,5,1,1
--   N=23 → 5,5,5,5,2,1
--   N=24 → 5,5,5,5,2,2
--   N=26 → 5,5,5,5,3,3
--   N=27 → 5,5,5,5,4,3
--   N=30 → 5,5,5,5,5,5 (exact fill)
-- And at smaller group counts (maxGroup = 2, i.e. a single 2-group
-- analysis): N=6 → 3,3; N=11 → impossible (>10); small raids just
-- balance because k is clamped to 0.
--
-- Side note on odds-mode splits: a compaction move from G5 → G4
-- crosses sides, so split balance can shift by one body when the
-- tail bridges sides. That's accepted — the user's explicit ask is
-- for visual density in the raid window; the balance panel updates
-- after the moves apply and the next pass re-ranks swaps against
-- the new state.
------------------------------------------------------------
local function Criterion_Compact(split, H, maxGroup, splitMode)
    if maxGroup < 1 then return end

    local perGroup = {}
    for g = 1, maxGroup do perGroup[g] = 0 end

    -- Aggregate both sides — compaction is raid-wide.
    local allPlayers = {}
    for _, p in ipairs(split.A.all) do
        if perGroup[p.group] ~= nil then
            perGroup[p.group] = perGroup[p.group] + 1
            table.insert(allPlayers, p)
        end
    end
    for _, p in ipairs(split.B.all) do
        if perGroup[p.group] ~= nil then
            perGroup[p.group] = perGroup[p.group] + 1
            table.insert(allPlayers, p)
        end
    end

    local N = #allPlayers
    if N == 0 then return end

    local k
    if N >= 5 * maxGroup then
        k = maxGroup
    else
        k = math.min(maxGroup - 2, math.floor(N / 5))
        if k < 0 then k = 0 end
    end

    local target = {}
    for g = 1, k do target[g] = 5 end
    local tailCount = maxGroup - k
    if tailCount > 0 then
        local remaining  = N - 5 * k
        local targetHigh = math.ceil(remaining / tailCount)
        local targetLow  = math.floor(remaining / tailCount)
        local numHighs   = remaining - targetLow * tailCount
        for g = k + 1, maxGroup do
            if (g - k) <= numHighs then
                target[g] = targetHigh
            else
                target[g] = targetLow
            end
        end
    end

    local over, under = {}, {}
    for g = 1, maxGroup do
        if perGroup[g] > target[g] then
            table.insert(over, g)
        elseif perGroup[g] < target[g] then
            table.insert(under, g)
        end
    end

    -- Walk the two pointer lists; move one body at a time until
    -- every group hits its target. Advance a pointer when its
    -- group is satisfied, bail out of a source if no unused player
    -- in it is available (earlier criteria may have claimed them).
    local oi, ui = 1, 1
    while oi <= #over and ui <= #under do
        local fromG = over[oi]
        local toG   = under[ui]
        if perGroup[fromG] <= target[fromG] then
            oi = oi + 1
        elseif perGroup[toG] >= target[toG] then
            ui = ui + 1
        else
            local moved = false
            for _, p in ipairs(allPlayers) do
                if p.group == fromG and not H.IsUsed(p.name) then
                    if H.AddMove(p, toG, "Consolidate") then
                        perGroup[fromG] = perGroup[fromG] - 1
                        perGroup[toG]   = perGroup[toG]   + 1
                        moved = true
                        break
                    end
                end
            end
            if not moved then oi = oi + 1 end
        end
    end
end

------------------------------------------------------------
-- Plan persistence helpers — the UI caches the previous plan
-- across refreshes so clicking one swap doesn't re-rank the
-- remaining steps. Serialize stores only names + target groups
-- (indices shift between refreshes); Reconcile re-attaches live
-- player refs and drops entries that are no longer needed.
------------------------------------------------------------
function SWAP:SerializePlan(suggestions)
    local out = {}
    for _, s in ipairs(suggestions or {}) do
        if s.op == "move" and s.p1 and s.targetGroup then
            table.insert(out, {
                op          = "move",
                p1Name      = s.p1.name,
                targetGroup = s.targetGroup,
                reason      = s.reason,
            })
        elseif s.p1 and s.p2 then
            table.insert(out, {
                op     = "swap",
                p1Name = s.p1.name,
                p2Name = s.p2.name,
                reason = s.reason,
            })
        end
    end
    return out
end

function SWAP:Reconcile(cachedPlan, roster, maxGroup, splitMode)
    local byName = {}
    for _, p in ipairs(roster) do byName[p.name] = p end

    local out = {}
    for _, e in ipairs(cachedPlan or {}) do
        if e.op == "move" then
            local p = byName[e.p1Name]
            if p and p.group ~= e.targetGroup then
                table.insert(out, {
                    op          = "move",
                    p1          = p,
                    targetGroup = e.targetGroup,
                    reason      = e.reason,
                })
            end
        elseif e.op == "swap" then
            local p1 = byName[e.p1Name]
            local p2 = byName[e.p2Name]
            if p1 and p2 then
                local p1A = IsSideA(p1.group, splitMode, maxGroup)
                local p2A = IsSideA(p2.group, splitMode, maxGroup)
                if p1A ~= p2A then
                    table.insert(out, {
                        op     = "swap",
                        p1     = p1,
                        p2     = p2,
                        reason = e.reason,
                    })
                end
            end
        end
    end
    return out
end

function SWAP:RemoveFromPlan(plan, op, p1Name, p2Name, targetGroup)
    if not plan then return false end
    for i, e in ipairs(plan) do
        if e.op == op then
            if op == "move" and e.p1Name == p1Name and e.targetGroup == targetGroup then
                table.remove(plan, i)
                return true
            end
            if op == "swap" then
                -- Swaps are symmetric; match either orientation.
                if (e.p1Name == p1Name and e.p2Name == p2Name) or
                   (e.p1Name == p2Name and e.p2Name == p1Name) then
                    table.remove(plan, i)
                    return true
                end
            end
        end
    end
    return false
end

------------------------------------------------------------
-- Entry point
------------------------------------------------------------
function SWAP:Suggest(roster, opts)
    opts = opts or {}
    local maxGroup  = opts.maxGroup  or 6
    local splitMode = opts.splitMode or "odds"

    local split = self:BuildSplit(roster, maxGroup, splitMode)
    local suggestions = {}
    local H = MakeHelpers(suggestions, split, maxGroup, splitMode)

    Criterion_Count(split, H)
    Criterion_Tanks(split, H)
    Criterion_Healers(split, H)
    Criterion_Lust(split, H)
    Criterion_Brez(split, H)
    Criterion_HealerSpecs(split, H)
    Criterion_MeleeRanged(split, H)
    Criterion_DebuffClasses(split, H)
    Criterion_PIPairs(split, H, opts.piAssignments, roster)
    -- Compact runs LAST so swap-driven group changes are already done
    -- and we only consolidate whatever the final split looks like.
    Criterion_Compact(split, H, maxGroup, splitMode)

    return {
        suggestions = suggestions,
        splitA      = split.A,
        splitB      = split.B,
        allPriests  = split.allPriests,
        mode        = splitMode,
        maxGroup    = maxGroup,
    }
end
