rSemiTruck = rSemiTruck or {}
rSemiTruck.ATATuning = rSemiTruck.ATATuning or {}
rSemiTruck.ATATuning.Init = rSemiTruck.ATATuning.Init or {}
rSemiTruck.ATATuning.InstallTest = rSemiTruck.ATATuning.InstallTest or {}
rSemiTruck.ATATuning.InstallComplete = rSemiTruck.ATATuning.InstallComplete or {}
rSemiTruck.ATATuning.UninstallTest = rSemiTruck.ATATuning.UninstallTest or {}
rSemiTruck.ATATuning.UninstallComplete = rSemiTruck.ATATuning.UninstallComplete or {}
rSemiTruck.ATATuning.Update = rSemiTruck.ATATuning.Update or {}

local middleWheelPartIds = {
    "TireMiddleLeft",
    "TireMiddleRight",
}

local bogieTirePairs = {
    { "TireMiddleLeft", "TireRearLeft" },
    { "TireMiddleRight", "TireRearRight" },
}

local tirePartToWheelId = {
    TireFrontLeft = "FrontLeft",
    TireFrontRight = "FrontRight",
    TireMiddleLeft = "RearLeft",
    TireMiddleRight = "RearLeft",
    TireRearLeft = "RearRight",
    TireRearRight = "RearRight",
}

local semiTruckScriptNames = {
    ["Base.SemiTruck"] = true,
    ["Base.SemiTruck_mil"] = true,
    ["Base.SemiTruckLite"] = true,
    ["Base.SemiTruckBox"] = true,
    ["Base.SemiTruckBox_mil"] = true,
}

local function getBaseATATuning()
    return rawget(_G, "ATATuning")
end

local function getATATuning2Utils()
    return rawget(_G, "ATATuning2Utils")
end

local function forEachLoadedVehicle(visitor)
    if type(visitor) ~= "function" then
        return
    end

    local cell = getCell and getCell() or nil
    if not cell or not cell.getVehicles then
        return
    end

    local vehicles = cell:getVehicles()
    if not vehicles then
        return
    end

    if type(vehicles.size) == "function" and type(vehicles.get) == "function" then
        for index = 0, vehicles:size() - 1 do
            visitor(vehicles:get(index), index)
        end
        return
    end

    if type(vehicles.iterator) == "function" then
        local ok, iterator = pcall(function() return vehicles:iterator() end)
        if ok and iterator then
            local index = 0
            while iterator:hasNext() do
                visitor(iterator:next(), index)
                index = index + 1
            end
            return
        end
    end

    local count = nil
    if type(vehicles.size) == "function" then
        local ok, sizeValue = pcall(function() return vehicles:size() end)
        if ok then
            count = sizeValue
        end
    elseif type(vehicles) == "table" then
        count = #vehicles
    end

    if count then
        for index = 0, count - 1 do
            local ok, vehicle = pcall(function() return vehicles[index] end)
            if not ok or vehicle == nil then
                ok, vehicle = pcall(function() return vehicles[index + 1] end)
            end
            if ok and vehicle then
                visitor(vehicle, index)
            end
        end
        return
    end

    if type(vehicles) == "table" then
        for index, vehicle in ipairs(vehicles) do
            visitor(vehicle, index - 1)
        end
    end
end

local function resolveWheelIndex(vehicle, part)
    if not vehicle or not part then
        return nil
    end

    local script = vehicle:getScript()
    local wheelIndex = part:getWheelIndex()

    if wheelIndex and wheelIndex >= 0 then
        if not script then
            return wheelIndex
        end

        local wheelCount = script:getWheelCount()
        if not wheelCount or wheelIndex < wheelCount then
            return wheelIndex
        end
    end

    if not script then
        return nil
    end

    local wheelCount = script:getWheelCount()
    if not wheelCount or wheelCount <= 0 then
        return nil
    end

    local partId = part:getId()
    local wheelId = tirePartToWheelId[partId]
    if not wheelId then
        wheelId = string.gsub(partId or "", "^Tire", "")
    end

    if wheelId == "" then
        return nil
    end

    for index = 0, wheelCount - 1 do
        local wheel = script:getWheel(index)
        if wheel and wheel:getId() == wheelId then
            return index
        end
    end

    return nil
end

local function setTireRemovedIfValid(vehicle, part, isRemoved)
    local wheelIndex = resolveWheelIndex(vehicle, part)
    if wheelIndex == nil then
        return false
    end

    vehicle:setTireRemoved(wheelIndex, isRemoved)
    return true
end

local function syncTireModels(vehicle, part, hasInflatedTire)
    if not part then
        return
    end

    part:setModelVisible("InflatedTirePlusWheel", hasInflatedTire)

    local utils = getATATuning2Utils()
    if utils and utils.ModelByItemName then
        if hasInflatedTire then
            utils.ModelByItemName(vehicle, part, part:getInventoryItem())
        else
            utils.ModelByItemName(vehicle, part)
        end
    end
end

local function isWSemiTruck(vehicle)
    if not vehicle then
        return false
    end

    local script = vehicle:getScript()
    if not script then
        return false
    end

    return semiTruckScriptNames[script:getFullName()] == true
end

local function updateMiddleWheelProtection(vehicle, hasProtection)
    if not isWSemiTruck(vehicle) then
        return
    end

    for _, partId in ipairs(middleWheelPartIds) do
        local wheelPart = vehicle:getPartById(partId)
        if wheelPart then
            wheelPart:setModelVisible("ATAProtection", hasProtection and wheelPart:getInventoryItem() ~= nil)
        end
    end
end

local function updateIndividualWheelProtection(part)
    if not part then
        return
    end

    local wheelPart = part:getParent()
    if not wheelPart then
        return
    end

    wheelPart:setModelVisible("ATAProtection", part:getInventoryItem() ~= nil and wheelPart:getInventoryItem() ~= nil)
end

local function getVehicleTuningModData(vehicle)
    if not vehicle then
        return nil
    end

    local modData = vehicle:getModData()
    modData.rSemiTruckATATuning = modData.rSemiTruckATATuning or {}
    modData.rSemiTruckATATuning.bogieTireState = modData.rSemiTruckATATuning.bogieTireState or {}
    return modData.rSemiTruckATATuning.bogieTireState
end

local function testMultiRequireDependencies(vehicle, keyvalues)
    if keyvalues.requireInstalled then
        local split = keyvalues.requireInstalled:split(";")
        for _, partId in ipairs(split) do
            local requiredPart = vehicle:getPartById(partId)
            if not requiredPart or not requiredPart:getInventoryItem() then
                return false
            end
        end
    end

    if keyvalues.requireUninstalled then
        local split = keyvalues.requireUninstalled:split(";")
        for _, partId in ipairs(split) do
            local blockedPart = vehicle:getPartById(partId)
            if blockedPart and blockedPart:getInventoryItem() then
                return false
            end
        end
    end

    return true
end

local function testMultiRequireCommon(vehicle, part, chr, keyvalues)
    if not VehicleUtils.testProfession(chr, keyvalues.professions) then
        return false
    end
    if not VehicleUtils.testRecipes(chr, keyvalues.recipes) then
        return false
    end
    if not VehicleUtils.testTraits(chr, keyvalues.traits) then
        return false
    end

    local typeToItem, tagToItem = VehicleUtils.getItems(chr:getPlayerNum())
    if not VehicleUtils.testItems(chr, keyvalues.items, typeToItem, tagToItem) then
        return false
    end

    if VehicleUtils.RequiredKeyNotFound(part, chr) then
        return false
    end

    return testMultiRequireDependencies(vehicle, keyvalues)
end

function rSemiTruck.ATATuning.InstallTest.IndividualWheelProtection(vehicle, part, chr)
    if ISVehicleMechanics.cheat then
        return true
    end

    local keyvalues = part:getTable("install")
    if not keyvalues then
        return false
    end
    if part:getInventoryItem() then
        return false
    end
    if not part:getItemType() or part:getItemType():isEmpty() then
        return false
    end

    return testMultiRequireCommon(vehicle, part, chr, keyvalues)
end

function rSemiTruck.ATATuning.UninstallTest.IndividualWheelProtection(vehicle, part, chr)
    if ISVehicleMechanics.cheat then
        return true
    end

    local keyvalues = part:getTable("uninstall")
    if not keyvalues then
        return false
    end
    if not part:getInventoryItem() then
        return false
    end
    if not part:getItemType() or part:getItemType():isEmpty() then
        return false
    end
    if not testMultiRequireCommon(vehicle, part, chr, keyvalues) then
        return false
    end
    if keyvalues.requireEmpty and round(part:getContainerContentAmount(), 3) > 0 then
        return false
    end

    local seatNumber = part:getContainerSeatNumber()
    local seatOccupied = (seatNumber ~= -1) and vehicle:isSeatOccupied(seatNumber)
    if keyvalues.requireEmpty and seatOccupied then
        return false
    end

    return true
end

local function getBogiePairState(vehicle, firstPartId, secondPartId)
    local stateByPair = getVehicleTuningModData(vehicle)
    if not stateByPair then
        return nil
    end

    local pairKey = firstPartId .. ":" .. secondPartId
    stateByPair[pairKey] = stateByPair[pairKey] or {}
    return stateByPair[pairKey]
end

local function setTireCondition(vehicle, part, condition)
    local clampedCondition = math.max(0, math.floor(condition or 0))
    if part:getCondition() ~= clampedCondition then
        part:setCondition(clampedCondition)
        vehicle:transmitPartCondition(part)
    end

    local modData = part:getModData()
    if modData then
        modData.atatuning = modData.atatuning or {}
        if modData.atatuning.health ~= clampedCondition then
            modData.atatuning.health = clampedCondition
            vehicle:transmitPartModData(part)
        end
    end

    return clampedCondition
end

local function syncBogieTirePair(vehicle, firstPartId, secondPartId)
    if not vehicle then
        return
    end

    local firstPart = vehicle:getPartById(firstPartId)
    local secondPart = vehicle:getPartById(secondPartId)
    if not firstPart or not secondPart then
        return
    end

    local pairState = getBogiePairState(vehicle, firstPartId, secondPartId)
    if not pairState then
        return
    end

    local currentFirst = firstPart:getCondition()
    local currentSecond = secondPart:getCondition()
    local hasFirstTire = firstPart:getInventoryItem() ~= nil
    local hasSecondTire = secondPart:getInventoryItem() ~= nil

    if pairState.firstCondition == nil or pairState.secondCondition == nil then
        pairState.firstCondition = currentFirst
        pairState.secondCondition = currentSecond
        setTireCondition(vehicle, firstPart, currentFirst)
        setTireCondition(vehicle, secondPart, currentSecond)
        return
    end

    local previousFirst = pairState.firstCondition
    local previousSecond = pairState.secondCondition
    local firstDamage = math.max(0, previousFirst - currentFirst)
    local secondDamage = math.max(0, previousSecond - currentSecond)
    local totalDamage = firstDamage + secondDamage

    local finalFirst = currentFirst
    local finalSecond = currentSecond

    if totalDamage > 0 then
        if hasFirstTire and hasSecondTire then
            local damageFirst = ZombRand(2) == 0
            if damageFirst then
                finalFirst = math.max(0, previousFirst - totalDamage)
                finalSecond = previousSecond
            else
                finalFirst = previousFirst
                finalSecond = math.max(0, previousSecond - totalDamage)
            end
        elseif hasFirstTire then
            finalFirst = math.max(0, previousFirst - totalDamage)
            finalSecond = currentSecond
        elseif hasSecondTire then
            finalFirst = currentFirst
            finalSecond = math.max(0, previousSecond - totalDamage)
        else
            finalFirst = currentFirst
            finalSecond = currentSecond
        end
    end

    if not hasFirstTire then
        finalFirst = currentFirst
    end
    if not hasSecondTire then
        finalSecond = currentSecond
    end

    pairState.firstCondition = setTireCondition(vehicle, firstPart, finalFirst)
    pairState.secondCondition = setTireCondition(vehicle, secondPart, finalSecond)
end

local function syncBogieTires(vehicle)
    if not vehicle then
        return
    end

    for _, pair in ipairs(bogieTirePairs) do
        syncBogieTirePair(vehicle, pair[1], pair[2])
    end
end

local function syncLoadedBogieState()
    forEachLoadedVehicle(function(vehicle)
        if vehicle and vehicle:getPartById("TireMiddleLeft") and vehicle:getPartById("TireRearLeft") then
            syncBogieTires(vehicle)
        end
    end)
end

local function consumeIndividualWheelProtectionDamage(vehicle, part)
    if not vehicle or not part then
        return
    end

    local item = part:getInventoryItem()
    if not item then
        return
    end

    local wheelPart = part:getParent()
    if not wheelPart or not wheelPart:getInventoryItem() then
        return
    end

    local wheelModData = wheelPart:getModData()
    wheelModData.atatuning = wheelModData.atatuning or {}
    if not wheelModData.atatuning.health then
        wheelModData.atatuning.health = wheelPart:getCondition()
    end

    local wheelCondition = wheelPart:getCondition()
    if wheelCondition < 100 then
        local absorbedDamage = 100 - wheelCondition
        if absorbedDamage > 0 then
            local newProtectionCondition = math.max(0, part:getCondition() - absorbedDamage)
            if newProtectionCondition ~= part:getCondition() then
                part:setCondition(newProtectionCondition)
                vehicle:transmitPartCondition(part)
            end

            wheelPart:setCondition(100)
            vehicle:transmitPartCondition(wheelPart)
        end
    end

    if string.match(wheelPart:getId(), "Tire") and wheelPart:getContainerContentAmount() < 10 then
        wheelPart:setContainerContentAmount(20, false, true)
    end

    vehicle:transmitPartModData(wheelPart)

    if part:getCondition() <= 0 then
        local areaCenter = vehicle:getAreaCenter(part:getArea())
        local square = nil
        if areaCenter then
            square = getCell():getGridSquare(areaCenter:getX(), areaCenter:getY(), vehicle:getZ())
        end
        part:setInventoryItem(nil)
        vehicle:transmitPartItem(part)
        if square then
            square:AddWorldInventoryItem(item, 0.5, 0.5, 0)
        end
        local tuning = getBaseATATuning()
        if tuning and tuning.UninstallComplete and tuning.UninstallComplete.CommonProtection then
            tuning.UninstallComplete.CommonProtection(vehicle, part, item)
        end
    end
end

function rSemiTruck.ATATuning.MiddleAxleWheelsProtection(vehicle, part)
    updateMiddleWheelProtection(vehicle, part and part:getInventoryItem() ~= nil)
end

function rSemiTruck.ATATuning.Init.MiddleAxleWheelsProtection(vehicle, part)
    local tuning = getBaseATATuning()
    if tuning and tuning.Init and tuning.Init.WheelsProtection then
        tuning.Init.WheelsProtection(vehicle, part)
    end

    rSemiTruck.ATATuning.MiddleAxleWheelsProtection(vehicle, part)
end

function rSemiTruck.ATATuning.InstallComplete.MiddleAxleWheelsProtection(vehicle, part)
    local tuning = getBaseATATuning()
    if tuning and tuning.InstallComplete and tuning.InstallComplete.WheelsProtection then
        tuning.InstallComplete.WheelsProtection(vehicle, part)
    end

    rSemiTruck.ATATuning.MiddleAxleWheelsProtection(vehicle, part)
end

function rSemiTruck.ATATuning.UninstallComplete.MiddleAxleWheelsProtection(vehicle, part, item)
    local tuning = getBaseATATuning()
    if tuning and tuning.UninstallComplete and tuning.UninstallComplete.WheelsProtection then
        tuning.UninstallComplete.WheelsProtection(vehicle, part, item)
    end

    rSemiTruck.ATATuning.MiddleAxleWheelsProtection(vehicle, part)
end

function rSemiTruck.ATATuning.Init.SafeTireNotAllModelsVisible(vehicle, part)
    local hasInflatedTire = part and part:getInventoryItem() ~= nil or false
    setTireRemovedIfValid(vehicle, part, not hasInflatedTire)
    syncTireModels(vehicle, part, hasInflatedTire)
end

function rSemiTruck.ATATuning.InstallComplete.SafeTireNotAllModelsVisible(vehicle, part)
    setTireRemovedIfValid(vehicle, part, false)
    syncTireModels(vehicle, part, true)
end

function rSemiTruck.ATATuning.UninstallComplete.SafeTireNotAllModelsVisible(vehicle, part, item)
    setTireRemovedIfValid(vehicle, part, true)
    syncTireModels(vehicle, part, false)
end

function rSemiTruck.ATATuning.Init.IndividualWheelProtection(vehicle, part)
    local tuning = getBaseATATuning()
    if tuning and tuning.Init and tuning.Init.Default then
        tuning.Init.Default(vehicle, part)
    end

    updateIndividualWheelProtection(part)
end

function rSemiTruck.ATATuning.InstallComplete.IndividualWheelProtection(vehicle, part)
    local tuning = getBaseATATuning()
    if tuning and tuning.InstallComplete and tuning.InstallComplete.CommonProtection then
        tuning.InstallComplete.CommonProtection(vehicle, part)
    end

    updateIndividualWheelProtection(part)
end

function rSemiTruck.ATATuning.UninstallComplete.IndividualWheelProtection(vehicle, part, item)
    local tuning = getBaseATATuning()
    if tuning and tuning.UninstallComplete and tuning.UninstallComplete.CommonProtection then
        tuning.UninstallComplete.CommonProtection(vehicle, part, item)
    end

    updateIndividualWheelProtection(part)
end

function rSemiTruck.ATATuning.Update.IndividualWheelProtection(vehicle, part, elapsedMinutes)
    consumeIndividualWheelProtectionDamage(vehicle, part)

    updateIndividualWheelProtection(part)
end

local bogieTireSyncTick = 0
Events.OnTick.Add(function()
    bogieTireSyncTick = (bogieTireSyncTick + 1) % 15
    if bogieTireSyncTick ~= 0 then
        return
    end

    syncLoadedBogieState()
end)