require "Vehicles/ISUI/ISVehicleMechanics"

local W900_MECHANICS_ORDER = {
	Base = true,
	["Base.SemiTruck"] = true,
	["Base.SemiTruck_mil"] = true,
	["Base.SemiTruckLite"] = true,
	["Base.SemiTruckBox"] = true,
	["Base.SemiTruckBox_mil"] = true,
    ["Base.SemiTrailerVan"] = true,
    ["Base.SemiTrailerVan_mil"] = true,
    ["Base.SemiTrailerVanCattle"] = true,
    ["Base.SemiTrailerLowBed"] = true,
    ["Base.SemiTrailerCistern"] = true,

}

local TIRE_PART_ORDER = {
	TireFrontLeft = 1,
	WheelChainFrontLeft = 2,
	TireFrontRight = 3,
	WheelChainFrontRight = 4,
	TireMiddleLeft = 5,
	WheelChainMiddleLeft = 6,
	TireMiddleRight = 7,
	WheelChainMiddleRight = 8,
	TireRearLeft = 9,
	WheelChainRearLeft = 10,
	TireRearRight = 11,
	WheelChainRearRight = 12,
}

local function isW900Vehicle(vehicle)
	if not vehicle or not vehicle.getScriptName then
		return false
	end

	return W900_MECHANICS_ORDER[vehicle:getScriptName()] == true
end

local function getPartOrder(partId, fallback)
	return TIRE_PART_ORDER[partId] or (1000 + fallback)
end

local function sortTireParts(parts)
	local indexedParts = {}
	for index, entry in ipairs(parts) do
		indexedParts[index] = {
			entry = entry,
			index = index,
			partId = entry.part and entry.part:getId() or "",
		}
	end

	table.sort(indexedParts, function(left, right)
		local leftOrder = getPartOrder(left.partId, left.index)
		local rightOrder = getPartOrder(right.partId, right.index)
		if leftOrder == rightOrder then
			return left.index < right.index
		end
		return leftOrder < rightOrder
	end)

	for index, item in ipairs(indexedParts) do
		parts[index] = item.entry
	end
	return indexedParts
end

local function reorderListboxTireCategory(listbox)
	if not listbox or not listbox.items then
		return
	end

	local selectedPartId
	local selectedRow = listbox.items[listbox.selected]
	if selectedRow and selectedRow.item and selectedRow.item.part then
		selectedPartId = selectedRow.item.part:getId()
	end

	local tireCategoryName = getText("IGUI_VehiclePartCattire")
	local headerIndex
	for index, row in ipairs(listbox.items) do
		if row.item and row.item.cat and row.text == tireCategoryName then
			headerIndex = index
			break
		end
	end

	if not headerIndex then
		return
	end

	local startIndex = headerIndex + 1
	local endIndex = #listbox.items
	for index = startIndex, #listbox.items do
		local row = listbox.items[index]
		if row and row.item and row.item.cat then
			endIndex = index - 1
			break
		end
	end

	if endIndex < startIndex then
		return
	end

	local slice = {}
	for index = startIndex, endIndex do
		local row = listbox.items[index]
		slice[#slice + 1] = {
			row = row,
			index = index,
			partId = row.item and row.item.part and row.item.part:getId() or "",
		}
	end

	table.sort(slice, function(left, right)
		local leftOrder = getPartOrder(left.partId, left.index)
		local rightOrder = getPartOrder(right.partId, right.index)
		if leftOrder == rightOrder then
			return left.index < right.index
		end
		return leftOrder < rightOrder
	end)

	for offset, item in ipairs(slice) do
		listbox.items[startIndex + offset - 1] = item.row
	end

	for index, row in ipairs(listbox.items) do
		row.itemindex = index
	end

	if selectedPartId then
		for index = startIndex, endIndex do
			local row = listbox.items[index]
			if row and row.item and row.item.part and row.item.part:getId() == selectedPartId then
				listbox.selected = index
				break
			end
		end
	end

	if listbox.selected > #listbox.items then
		listbox.selected = #listbox.items
	end
	if listbox.selected and listbox.selected > 0 then
		listbox:ensureVisible(listbox.selected)
	end
end

local originalInitParts = ISVehicleMechanics.initParts

function ISVehicleMechanics:initParts()
	originalInitParts(self)

	if not isW900Vehicle(self.vehicle) then
		return
	end

	local tireCategory = self.vehiclePart and self.vehiclePart.tire
	if tireCategory and tireCategory.parts then
		sortTireParts(tireCategory.parts)
	end

	reorderListboxTireCategory(self.listbox)
end