-- Arena nameplates.
-- Lets you replace enemy names in arena with arena# / spec / both.

local AC = _G.ArenaCore
if not AC then return end

local RefreshArenaPlates
local GetTransformedNameText
local EnsureNameOverrideHook
local TransformNameplateText

local arenaPlates = {} -- [1..3] = nameplate frame or nil
local nameplateToArenaIndex = {} -- [nameplate] = 1/2/3
local arenaIndexToClass = {} -- [1..3] = "WARRIOR", "MAGE", etc.
local arenaIndexToSpec = {} -- [1..3] = specID

-- Spec ID -> display names (used by RefreshArenaPlates).
local SPEC_NAMES_FULL = {
    -- Death Knight
    [250] = "Blood", [251] = "Frost", [252] = "Unholy",
    -- Demon Hunter
    [577] = "Havoc", [581] = "Vengeance",
    -- Druid
    [102] = "Balance", [103] = "Feral", [104] = "Guardian", [105] = "Restoration",
    -- Evoker
    [1467] = "Devastation", [1468] = "Preservation", [1473] = "Augmentation",
    -- Hunter
    [253] = "Beast Mastery", [254] = "Marksmanship", [255] = "Survival",
    -- Mage
    [62] = "Arcane", [63] = "Fire", [64] = "Frost",
    -- Monk
    [268] = "Brewmaster", [270] = "Mistweaver", [269] = "Windwalker",
    -- Paladin
    [65] = "Holy", [66] = "Protection", [70] = "Retribution",
    -- Priest
    [256] = "Discipline", [257] = "Holy", [258] = "Shadow",
    -- Rogue
    [259] = "Assassination", [260] = "Outlaw", [261] = "Subtlety",
    -- Shaman
    [262] = "Elemental", [263] = "Enhancement", [264] = "Restoration",
    -- Warlock
    [265] = "Affliction", [266] = "Demonology", [267] = "Destruction",
    -- Warrior
    [71] = "Arms", [72] = "Fury", [73] = "Protection",
}

local SPEC_NAMES_SHORT = {
    -- Death Knight
    [250] = "Blood", [251] = "Frost", [252] = "Unholy",
    -- Demon Hunter
    [577] = "Havoc", [581] = "Veng",
    -- Druid
    [102] = "Balance", [103] = "Feral", [104] = "Guardian", [105] = "Resto",
    -- Evoker
    [1467] = "Dev", [1468] = "Pres", [1473] = "Aug",
    -- Hunter
    [253] = "BM", [254] = "MM", [255] = "Survival",
    -- Mage
    [62] = "Arcane", [63] = "Fire", [64] = "Frost",
    -- Monk
    [268] = "Brew", [270] = "MW", [269] = "WW",
    -- Paladin
    [65] = "Holy", [66] = "Prot", [70] = "Ret",
    -- Priest
    [256] = "Disc", [257] = "Holy", [258] = "Shadow",
    -- Rogue
    [259] = "Assa", [260] = "Outlaw", [261] = "Sub",
    -- Shaman
    [262] = "Ele", [263] = "Enh", [264] = "Resto",
    -- Warlock
    [265] = "Aff", [266] = "Demo", [267] = "Destro",
    -- Warrior
    [71] = "Arms", [72] = "Fury", [73] = "Prot",
}

-- Nameplate unit tokens are safe for direct Unit* checks.
local function IsEnemyPlayer(unit)
    if not unit then return false end
    return UnitIsPlayer(unit) and UnitIsEnemy("player", unit)
end

local function EnsureNameOverrideHook(unitFrame)
    if not unitFrame or unitFrame.__acArenaNameOverrideHooked then return end

    unitFrame.__acArenaNameOverrideHooked = true
    unitFrame.__acOriginalUpdateNameOverride = unitFrame.UpdateNameOverride

    -- Don't overwrite Blizzard mixins (taint bait). Post-hook instead.
    if type(unitFrame.UpdateNameOverride) == "function" then
        hooksecurefunc(unitFrame, "UpdateNameOverride", function(self)
            local db = AC.DB and AC.DB.profile and AC.DB.profile.arenaNameplates
            if not db or not db.enabled then return end
            if not IsActiveBattlefieldArena() then return end

            local unit = self.unit
            if not unit or not unit:match("^nameplate") then return end

            local arenaIndex = self.arenaID
            if not arenaIndex then return end

            local newText = GetTransformedNameText(arenaIndex, db.mode, db.useShortNames)
            if not newText then return end

            if self.__acArenaNameText and self.__acArenaNameText.SetText then
                self.__acArenaNameText:SetText(newText)
                self.__acArenaNameText:Show()
            end

            if self.name then
                self.name:SetText("")
                self.name:SetAlpha(0)
            end
        end)
    end
end

RefreshArenaPlates = function(shouldWipe)
    if shouldWipe then
        wipe(nameplateToArenaIndex)
        wipe(arenaIndexToClass)
        wipe(arenaIndexToSpec)
    end

    -- Cache arena class/spec info from arena tokens.
    for i = 1, 3 do
        local specID = GetArenaOpponentSpec(i)
        if specID then
            arenaIndexToSpec[i] = specID
        end
    end

    local nameplates = C_NamePlate.GetNamePlates()
    if not nameplates then return end

    local foundArenas = {}
    local allEnemyPlates = {}
    local plateToClass = {}
    local plateToSpec = {}

    -- Pass 1: validate any existing arenaID tags on enemy player nameplates.
    local alreadyTaggedCount = 0
    for _, plate in ipairs(nameplates) do
        local uf = plate.UnitFrame or plate.unitFrame
        if uf and uf.arenaID then
            local unit = uf.unit
            if unit and IsEnemyPlayer(unit) then
                local cachedIndex = nameplateToArenaIndex[plate]
                if cachedIndex == uf.arenaID then
                    alreadyTaggedCount = alreadyTaggedCount + 1
                    foundArenas[uf.arenaID] = plate
                else
                    uf.arenaID = nil
                end
            else
                uf.arenaID = nil
            end
        end
    end

    -- If all 3 are tagged and still valid, we're done.
    if alreadyTaggedCount == 3 and foundArenas[1] and foundArenas[2] and foundArenas[3] then
        for i = 1, 3 do arenaPlates[i] = foundArenas[i] end
        return
    end

    -- Pass 2: collect enemy players and try indirect matching via target/focus/mouseover.
    -- Can't do UnitIsUnit(nameplate, "arena1") directly (secret boolean),
    -- but we can go through intermediates.
    for _, plate in ipairs(nameplates) do
        local uf = plate.UnitFrame or plate.unitFrame
        if uf then
            local unit = uf.unit
            if unit and IsEnemyPlayer(unit) then
                table.insert(allEnemyPlates, plate)
                local _, class = UnitClass(unit)
                plateToClass[plate] = class

                -- Validate cached index against class changes.
                local cachedIndex = nameplateToArenaIndex[plate]
                if cachedIndex then
                    local cachedClass = arenaIndexToClass[cachedIndex]
                    if cachedClass and cachedClass ~= class then
                        nameplateToArenaIndex[plate] = nil
                        uf.arenaID = nil
                        cachedIndex = nil
                    end
                end

                -- Indirect matching.
                for i = 1, 3 do
                    if not foundArenas[i] then
                        local arenaUnit = "arena" .. i
                        if (UnitIsUnit(unit, "target") and UnitIsUnit("target", arenaUnit)) or
                           (UnitIsUnit(unit, "focus") and UnitIsUnit("focus", arenaUnit)) or
                           (UnitIsUnit(unit, "mouseover") and UnitIsUnit("mouseover", arenaUnit)) then
                            nameplateToArenaIndex[plate] = i
                            arenaIndexToClass[i] = class
                            local specID = GetArenaOpponentSpec(i)
                            if specID then
                                arenaIndexToSpec[i] = specID
                                plateToSpec[plate] = specID
                            end
                            uf.arenaID = i
                            foundArenas[i] = plate
                        end
                    end
                end
            else
                -- Not an enemy player, clear any stale tags.
                if nameplateToArenaIndex[plate] then
                    nameplateToArenaIndex[plate] = nil
                end
                if uf and uf.arenaID then
                    uf.arenaID = nil
                end
            end
        end
    end

    -- Propagate cached spec info.
    for _, plate in ipairs(allEnemyPlates) do
        if not plateToSpec[plate] then
            local cachedIndex = nameplateToArenaIndex[plate]
            if cachedIndex and arenaIndexToSpec[cachedIndex] then
                plateToSpec[plate] = arenaIndexToSpec[cachedIndex]
            end
        end
    end

    -- Pass 3: class/spec deduction for unmatched plates.
    if #allEnemyPlates == 3 then
        for _, plate in ipairs(allEnemyPlates) do
            if not nameplateToArenaIndex[plate] then
                local class = plateToClass[plate]
                if class then
                    for i = 1, 3 do
                        if arenaIndexToClass[i] == class and not foundArenas[i] then
                            local cachedSpec = arenaIndexToSpec[i]
                            local pSpec = plateToSpec[plate]
                            if cachedSpec and pSpec then
                                if cachedSpec == pSpec then
                                    nameplateToArenaIndex[plate] = i
                                    local uf = plate.UnitFrame or plate.unitFrame
                                    if uf then uf.arenaID = i end
                                    foundArenas[i] = plate
                                    break
                                end
                            else
                                local classCount = 0
                                for _, p in ipairs(allEnemyPlates) do
                                    if plateToClass[p] == class and not nameplateToArenaIndex[p] then
                                        classCount = classCount + 1
                                    end
                                end
                                if classCount == 1 then
                                    nameplateToArenaIndex[plate] = i
                                    local uf = plate.UnitFrame or plate.unitFrame
                                    if uf then uf.arenaID = i end
                                    foundArenas[i] = plate
                                    break
                                end
                            end
                        end
                    end
                end
            end
        end
    end

    -- Pass 4: process of elimination.
    local foundCount = 0
    local missingIndex = nil
    for i = 1, 3 do
        if foundArenas[i] then
            foundCount = foundCount + 1
        else
            missingIndex = i
        end
    end

    if foundCount == 2 and missingIndex and #allEnemyPlates == 3 then
        for _, plate in ipairs(allEnemyPlates) do
            local isAlreadyFound = false
            for _, foundPlate in pairs(foundArenas) do
                if foundPlate == plate then
                    isAlreadyFound = true
                    break
                end
            end
            if not isAlreadyFound then
                nameplateToArenaIndex[plate] = missingIndex
                arenaIndexToClass[missingIndex] = plateToClass[plate]
                local specID = GetArenaOpponentSpec(missingIndex)
                if specID then
                    arenaIndexToSpec[missingIndex] = specID
                end
                local uf = plate.UnitFrame or plate.unitFrame
                if uf then uf.arenaID = missingIndex end
                break
            end
        end
    end

    -- Populate arenaPlates from results.
    for i = 1, 3 do
        arenaPlates[i] = foundArenas[i]
    end
end

-- ============================================================================
-- PLATYNATOR
-- ============================================================================
-- If Platynator is around, use its API for text overrides.

local function IsPlatynatorAPIAvailable()
    return Platynator and Platynator.API and Platynator.API.SetUnitTextOverride
end

local function GetSpecName(arenaIndex, useShortNames)
    local specID = arenaIndexToSpec[arenaIndex]
    local specTable = useShortNames and SPEC_NAMES_SHORT or SPEC_NAMES_FULL
    
    -- Try local mapping first.
    local specName = specID and specTable[specID]
    
    -- Fallback: ask Blizzard for the spec name.
    if not specName and specID and GetSpecializationInfoByID then
        local success, name = pcall(GetSpecializationInfoByID, specID)
        if success and type(name) == "string" and name ~= "" then
            specName = name
        end
    end
    
    -- Fallback to class name.
    if not specName then
        local className = arenaIndexToClass[arenaIndex]
        return className or "Unknown"
    end
    
    return specName
end

GetTransformedNameText = function(arenaIndex, mode, useShortNames)
    if not arenaIndex or arenaIndex < 1 or arenaIndex > 3 then return nil end
    
    if mode == "arena" then
        return tostring(arenaIndex)
    elseif mode == "spec" then
        return GetSpecName(arenaIndex, useShortNames)
    elseif mode == "both" then
        local specName = GetSpecName(arenaIndex, useShortNames)
        return specName .. " " .. arenaIndex
    end
    
    return nil
end

-- ============================================================================
-- PLATYNATOR OVERRIDES
-- ============================================================================

local function UpdatePlatynatorOverrides()
    if not IsPlatynatorAPIAvailable() then return end
    
    local db = AC.DB and AC.DB.profile and AC.DB.profile.arenaNameplates
    if not db or not db.enabled then return end
    
    RefreshArenaPlates()
    
    for i = 1, 3 do
        local nameplate = arenaPlates[i]
        if nameplate then
            -- Get the nameplate unit token (nameplate1, nameplate2, etc.).
            local unit = nameplate.namePlateUnitToken
            if unit and IsValidArenaPlayer(unit) then
                local transformedName = GetTransformedNameText(i, db.mode, db.useShortNames)
                if transformedName then
                    Platynator.API.SetUnitTextOverride(unit, transformedName, nil)
                end
            end
        end
    end
end

local function CreateCustomTextElements(unitFrame)
    -- Create custom text elements if they don't exist.
    if not unitFrame.__acArenaNameText then
        unitFrame.__acArenaNameText = unitFrame:CreateFontString(nil, "OVERLAY")
        unitFrame.__acArenaNameText:SetFont("Interface\\AddOns\\ArenaCore\\Media\\Fonts\\arenacore.ttf", 12, "OUTLINE")
        unitFrame.__acArenaNameText:SetIgnoreParentScale(true)
        unitFrame.__acArenaNameText:Hide()
    end
end

TransformNameplateText = function(nameplate, unitFrame, arenaIndex, mode, useShortNames, useClassColors)
    if not unitFrame or not arenaIndex then return end
    
    -- Name text element.
    local nameText = unitFrame.name
    if not nameText then return end
    
    -- Health bar for positioning.
    local healthBar = unitFrame.healthBar
    if not healthBar then return end
    
    -- Make sure our custom fontstring exists.
    CreateCustomTextElements(unitFrame)
    
    -- Pick the text.
    local newText
    if mode == "arena" then
        newText = tostring(arenaIndex)
    elseif mode == "spec" then
        newText = GetSpecName(arenaIndex, useShortNames)
    elseif mode == "both" then
        local specName = GetSpecName(arenaIndex, useShortNames)
        newText = specName .. " " .. arenaIndex
    else
        -- Default mode: restore original name and hide custom text.
        if unitFrame.__acArenaNameText then
            unitFrame.__acArenaNameText:Hide()
            unitFrame.__acArenaNameText:SetText("")
        end
        if nameText then
            nameText:SetAlpha(1)
            if CompactUnitFrame_UpdateName then
                pcall(CompactUnitFrame_UpdateName, unitFrame)
            end
        end
        return
    end
    
    -- Text color: class if enabled, otherwise white.
    local r, g, b
    if useClassColors then
        -- Use cached class (avoid extra token API calls).
        local class = arenaIndexToClass[arenaIndex]
        if class and RAID_CLASS_COLORS[class] then
            local classColor = RAID_CLASS_COLORS[class]
            r, g, b = classColor.r, classColor.g, classColor.b
        else
            r, g, b = 1, 1, 1
        end
    else
        r, g, b = 1, 1, 1
    end
    
    -- Make sure our post-hook exists so updates don't revert visuals.
    EnsureNameOverrideHook(unitFrame)

    if unitFrame.__acArenaNameText then
        unitFrame.__acArenaNameText:ClearAllPoints()
        if healthBar then
            unitFrame.__acArenaNameText:SetPoint("CENTER", nameText or healthBar, "CENTER", 0, 0)
        else
            unitFrame.__acArenaNameText:SetPoint("CENTER", unitFrame, "CENTER", 0, 0)
        end
        unitFrame.__acArenaNameText:SetText(newText)
        unitFrame.__acArenaNameText:SetTextColor(r, g, b, 1)
        unitFrame.__acArenaNameText:Show()
    end

    if nameText then
        nameText:SetText("")
        nameText:SetAlpha(0)
    end
end

-- ============================================================================
-- UNIT VALIDATION
-- ============================================================================

local function IsValidArenaPlayer(unit)
    if not unit then return false end
    return IsEnemyPlayer(unit)
end

-- ============================================================================
-- UPDATE
-- ============================================================================

local function UpdateArenaNameplate(nameplate, unitFrame, arenaIndex)
    if not unitFrame or not arenaIndex then return end
    
    -- Settings.
    local db = AC.DB and AC.DB.profile and AC.DB.profile.arenaNameplates
    if not db or not db.enabled then return end
    
    TransformNameplateText(nameplate, unitFrame, arenaIndex, db.mode, db.useShortNames, db.useClassColors)
end

local function RestoreOriginalName(unitFrame)
    if not unitFrame then return end

    if unitFrame.name then
        unitFrame.name:SetAlpha(1)
        if CompactUnitFrame_UpdateName then
            pcall(CompactUnitFrame_UpdateName, unitFrame)
        end
    end

    if unitFrame.__acArenaNameText then
        unitFrame.__acArenaNameText:Hide()
        unitFrame.__acArenaNameText:SetText("")
    end
end

local function RestoreAllOriginalNames()
    -- Only touch nameplates we actually modified.
    local plates = C_NamePlate.GetNamePlates()
    if not plates then return end
    for _, namePlate in ipairs(plates) do
        local unitFrame = namePlate.unitFrame or namePlate.UnitFrame
        if unitFrame and unitFrame.__acArenaNameText then
            RestoreOriginalName(unitFrame)
        end
    end
end

local function UpdateAllArenaNameplates()
    -- Refresh the arena plates cache first.
    RefreshArenaPlates()
    
    -- Iterate through cached arena nameplates.
    for i = 1, 3 do
        local nameplate = arenaPlates[i]
        if nameplate then
            local unitFrame = nameplate.unitFrame or nameplate.UnitFrame
            if unitFrame then
                UpdateArenaNameplate(nameplate, unitFrame, i)
            end
        end
    end
end

-- ============================================================================
-- EVENTS
-- ============================================================================

local eventFrame = CreateFrame("Frame")

local function RegisterEvents()
    eventFrame:RegisterEvent("PLAYER_ENTERING_WORLD")
    eventFrame:RegisterEvent("ARENA_OPPONENT_UPDATE")
    eventFrame:RegisterEvent("ARENA_PREP_OPPONENT_SPECIALIZATIONS")
    eventFrame:RegisterEvent("NAME_PLATE_UNIT_ADDED")
    eventFrame:RegisterEvent("NAME_PLATE_UNIT_REMOVED")
    eventFrame:RegisterEvent("PLAYER_TARGET_CHANGED")
    eventFrame:RegisterEvent("PLAYER_FOCUS_CHANGED")
    eventFrame:RegisterEvent("UPDATE_MOUSEOVER_UNIT")
    eventFrame:RegisterEvent("GROUP_ROSTER_UPDATE")
end

if AC.IsMidnight and AC:IsMidnight() then
    C_Timer.After(0.1, RegisterEvents)
else
    RegisterEvents()
end

eventFrame:SetScript("OnEvent", function(self, event, ...)
    local db = AC.DB and AC.DB.profile and AC.DB.profile.arenaNameplates
    
    -- Leaving arena or feature disabled.
    if not IsActiveBattlefieldArena() or not db or not db.enabled then
        -- This is kind of important: always restore nameplates when bailing.
        RestoreAllOriginalNames()
        
        -- Clear cache to prevent stale data.
        wipe(arenaPlates)
        return
    end
    
    if event == "PLAYER_ENTERING_WORLD" or event == "ARENA_OPPONENT_UPDATE" or event == "ARENA_PREP_OPPONENT_SPECIALIZATIONS" then
        RefreshArenaPlates(true)
        C_Timer.After(0.1, UpdateAllArenaNameplates)
        C_Timer.After(0.15, UpdatePlatynatorOverrides)
    elseif event == "GROUP_ROSTER_UPDATE" then
        -- Solo Shuffle round transitions: wipe and refresh.
        RestoreAllOriginalNames()
        RefreshArenaPlates(true)
        C_Timer.After(0.2, UpdateAllArenaNameplates)
    elseif event == "PLAYER_TARGET_CHANGED" or event == "PLAYER_FOCUS_CHANGED" or event == "UPDATE_MOUSEOVER_UNIT" then
        RefreshArenaPlates()
        C_Timer.After(0.05, UpdateAllArenaNameplates)
    elseif event == "NAME_PLATE_UNIT_ADDED" then
        local unit = ...
        local nameplate = C_NamePlate.GetNamePlateForUnit(unit)
        if nameplate then
            local unitFrame = nameplate.unitFrame or nameplate.UnitFrame
            if unitFrame then
                RestoreOriginalName(unitFrame)
                if IsEnemyPlayer(unit) then
                    CreateCustomTextElements(unitFrame)
                end
            end
        end
        RefreshArenaPlates()
        C_Timer.After(0.05, UpdateAllArenaNameplates)
    elseif event == "NAME_PLATE_UNIT_REMOVED" then
        local unit = ...
        local nameplate = C_NamePlate.GetNamePlateForUnit(unit)
        if nameplate then
            local unitFrame = nameplate.unitFrame or nameplate.UnitFrame
            if unitFrame then
                RestoreOriginalName(unitFrame)
                unitFrame.arenaID = nil
            end
            nameplateToArenaIndex[nameplate] = nil
        end
        RefreshArenaPlates()
    end
end)

-- ============================================================================
-- PUBLIC API
-- ============================================================================

function AC:RefreshArenaNameplates()
    UpdateAllArenaNameplates()
end

-- Initialize on load.
C_Timer.After(1, function()
    if IsActiveBattlefieldArena() then
        UpdateAllArenaNameplates()
    end
end)
