--@include /koptilnya/libs/constants.txt
--@include ./powertrain_component.txt

---@type PowertrainComponent
local PowertrainComponent = require('./powertrain_component.txt')

require('/koptilnya/libs/constants.txt')

---@alias DifferentialSplitStrategyFn fun(tq: number, aW: number, bW: number, aI: number, bI: number, biasAB: number, preload: number, stiffness: number, powerRamp: number, coastRamp: number, slipTorque: number): number, number

---@class DifferentialConfig: PowertrainComponentConfig
---@field Type? string
---@field FinalDrive? number
---@field Bias? number
---@field CoastRamp? number
---@field PowerRamp? number
---@field Preload? number
---@field Stiffness? number
---@field SlipTorque? number
---@field SteerLock? number
---@field ToeAngle? number

---@class Differential: PowertrainComponent
---@field finalDrive number
---@field biasAB number
---@field coastRamp number
---@field powerRamp number
---@field preload number
---@field stiffness number
---@field slipTorque number
---@field private splitStrategy DifferentialSplitStrategyFn
---@field steerLock number
---@field toeAngle number
---@field steerAngle number
---@field private mzDiff number
local Differential = class('Differential', PowertrainComponent)

---@private
---@param vehicle KPTLVehicle
---@param name string
---@param config DifferentialConfig
function Differential:initialize(vehicle, name, config)
	PowertrainComponent.initialize(self, vehicle, name, config)

	if CLIENT then
		return
	end

	self.wireOutputs = {
		[string.format('%s_Torque', self.name)] = 'number',
		[string.format('%s_RPMDiff', self.name)] = 'number',
	}

	self.finalDrive = config.FinalDrive or 4
	self.inertia = config.Inertia or 0.2
	self.biasAB = config.Bias or 0.5
	self.coastRamp = config.CoastRamp or 0.5
	self.powerRamp = config.PowerRamp or 1
	self.preload = config.Preload or 10
	self.stiffness = config.Stiffness or 0.1
	self.slipTorque = config.SlipTorque or 1000
	self.splitStrategy = Differential.getSplitStrategy(config.Type or DIFFERENTIAL_TYPES.Open)

	self.steerLock = config.SteerLock or 0
	self.toeAngle = config.ToeAngle or 0
	self.steerAngle = 0

	self.mzDiff = 0
	self.rpmDiff = 0
end

---@param component PowertrainComponent
---@return nil
function Differential:linkComponent(component)
	---@diagnostic disable-next-line: undefined-field
	if not component:isInstanceOf(PowertrainComponent) then
		return
	end

	if self.outputA == nil then
		self.outputA = component
		component.input = self
	elseif self.outputB == nil then
		self.outputB = component
		component.input = self
	end
end

---@return nil
function Differential:updateWireOutputs()
	PowertrainComponent.updateWireOutputs(self)

	wire.ports[string.format('%s_Torque', self.name)] = self.torque
	wire.ports[string.format('%s_RPMDiff', self.name)] = self.rpmDiff

	if self.outputA ~= nil then
		self.outputA:updateWireOutputs()
	end

	if self.outputB ~= nil then
		self.outputB:updateWireOutputs()
	end
end

---@return number
function Differential:queryInertia()
	if self.outputA == nil or self.outputB == nil then
		return self.inertia
	end

	return self.inertia + (self.outputA:queryInertia() + self.outputB:queryInertia()) / math.pow(self.finalDrive, 2)
end

---@return number
function Differential:queryAngularVelocity(angularVelocity)
	self.angularVelocity = angularVelocity

	if self.outputA == nil or self.outputB == nil then
		return angularVelocity
	end

	local aW = self.outputA:queryAngularVelocity(angularVelocity)
	local bW = self.outputB:queryAngularVelocity(angularVelocity)

	return (aW + bW) * self.finalDrive * 0.5
end

---@param torque number
---@param inertia number
---@return number
function Differential:forwardStep(torque, inertia)
	if self.outputA == nil or self.outputB == nil then
		return torque
	end

	local aW = self.outputA:queryAngularVelocity(self.angularVelocity)
	local bW = self.outputB:queryAngularVelocity(self.angularVelocity)
	local aI = self.outputA:queryInertia()
	local bI = self.outputB:queryInertia()

	self.torque = torque * self.finalDrive

	local tqA, tqB = self.splitStrategy(
		self.torque,
		aW,
		bW,
		aI,
		bI,
		self.biasAB,
		self.preload,
		self.stiffness,
		self.powerRamp,
		self.coastRamp,
		self.slipTorque
	)

	tqA = self.outputA:forwardStep(tqA, inertia * 0.5 * math.pow(self.finalDrive, 2) + aI)
	tqB = self.outputB:forwardStep(tqB, inertia * 0.5 * math.pow(self.finalDrive, 2) + bI)

	if self.CONFIG.Type == DIFFERENTIAL_TYPES.Locking then
		local avg = (self.outputA.angularVelocity + self.outputB.angularVelocity) / 2

		self.outputA:queryAngularVelocity(avg)
		self.outputB:queryAngularVelocity(avg)
	end

	self.rpmDiff = self.outputA:getRPM() - self.outputB:getRPM()

	-- // REFACTOR
	if self.outputA.customWheel ~= nil and self.outputB.customWheel ~= nil then
		if self.steerLock ~= 0 then
			local outputA = self.outputA --[[@as Wheel]]
			local outputB = self.outputB --[[@as Wheel]]

			local steerInertia = (aI + bI)
			local driverInput = 3

			local localVelocityLength = chip():getVelocity():getLength()
			local MPH = localVelocityLength * (15 / 352)
			local KPH = MPH * 1.609

			local assist = math.clamp(10.0 / math.sqrt(KPH / 3), 2.0, 10.0) / 2

			local inputForce = driverInput * assist

			local inputTorque = self.vehicle.steer * inputForce

			local mzDiff = (outputB.customWheel.mz - outputA.customWheel.mz) * TICK_INTERVAL / 2

			self.mzDiff = math.lerp(0.1, self.mzDiff, mzDiff)

			local steerTorque = self.mzDiff * 1 + inputTorque

			local steerAngularAccel = steerTorque / steerInertia

			self.steerAngle = math.clamp(self.steerAngle + math.deg(steerAngularAccel) * TICK_INTERVAL, -self.steerLock, self.steerLock)

			local wheelbase = 2.05
			local trackWidth = 1.124
			local radius = wheelbase / math.tan(math.rad(self.steerAngle))

			local innerAngle = math.deg(math.atan(wheelbase / (radius - (trackWidth / 2))))
			local outerAngle = math.deg(math.atan(wheelbase / (radius + (trackWidth / 2))))

			outputA.customWheel.steerAngle = outerAngle + self.toeAngle * outputA.customWheel.direction
			outputB.customWheel.steerAngle = innerAngle + self.toeAngle * outputB.customWheel.direction
		else
			local outputA = self.outputA --[[@as Wheel]]
			local outputB = self.outputB --[[@as Wheel]]

			outputA.customWheel.steerAngle = self.toeAngle * outputA.customWheel.direction
			outputB.customWheel.steerAngle = self.toeAngle * outputB.customWheel.direction
		end
	end

	return tqA + tqB
end

---@return number, number
local function _openDiffTorqueSplit(tq, aW, bW, aI, bI, biasAB, preload, stiffness, powerRamp, coastRamp, slipTorque)
	return tq * (1 - biasAB), tq * biasAB
end

---@return number, number
local function _lockingDiffTorqueSplit(tq, aW, bW, aI, bI, biasAB, preload, stiffness, powerRamp, coastRamp, slipTorque)
	return tq * 0.5, tq * 0.5
	-- -- Разность скоростей между выходами
	-- local deltaW = aW - bW

	-- -- Расчет "проскальзывающего" момента между сторонами
	-- local slip = math.abs(deltaW)

	-- -- Применяем преднагрузку и жесткость, чтобы определить базовый момент блокировки
	-- local lockingTorque = preload + stiffness * deltaW

	-- -- Определение направления нагрузки: ускорение или торможение
	-- local ramp = tq >= 0 and powerRamp or coastRamp

	-- -- Применяем ramp к разности скоростей, усиливая/ослабляя блокировку
	-- local rampTorque = slip * ramp

	-- -- Итоговый момент передачи между сторонами
	-- local transferTorque = lockingTorque + rampTorque

	-- -- Ограничиваем передачу момента порогом проскальзывания
	-- if transferTorque > slipTorque then
	-- 	transferTorque = slipTorque
	-- end

	-- -- Учитываем биас (смещение баланса)
	-- local maxBias = biasAB
	-- local biasLimit = tq * (maxBias - 1) / (maxBias + 1)

	-- -- Финальный момент, передаваемый между сторонами, с учетом направления
	-- local direction = deltaW > 0 and -1 or 1
	-- local finalTorque = direction * transferTorque

	-- -- Применяем финальный момент к каждой стороне, ограничивая его bias'ом
	-- local aTorque = tq * 0.5 + math.min(finalTorque, biasLimit)
	-- local bTorque = tq * 0.5 - math.min(finalTorque, biasLimit)

	-- -- Возвращаем обновленные значения угловых скоростей
	-- return aTorque, bTorque

	-- local aTq = tq * 0.5
	-- local bTq = tq * 0.5

	-- local syncTorque = (aW - bW) * 1e6 * (aI + bI) / TICK_INTERVAL

	-- aTq = aTq - syncTorque
	-- bTq = bTq + syncTorque

	-- -- print(aTq, bTq)

	-- return aTq, bTq

	-- local sumI = aI + bI
	-- local w = aI / sumI * aW + bI / sumI * bW
	-- local aTqCorr = (w - aW) * aI / TICK_INTERVAL
	-- aTqCorr = aTqCorr * stiffness

	-- local bTqCorr = (w - bW) * bI / TICK_INTERVAL
	-- bTqCorr = bTqCorr * stiffness

	-- local biasA = math.clamp(0.5 + (bW - aW) * 0.1 * stiffness, 0, 1)

	-- return tq * biasA + aTqCorr, tq * (1 - biasA) * bTqCorr
end

---@return number, number
local function _VLSDTorqueSplit(tq, aW, bW, aI, bI, biasAB, preload, stiffness, powerRamp, coastRamp, slipTorque)
	if aW < 0 or bW < 0 then
		return tq * (1 - biasAB), tq * biasAB
	end

	local c = tq > 0 and powerRamp or coastRamp
	local totalW = math.abs(aW) + math.abs(bW)
	local slip = 0
	if aW > 0 or bW > 0 then
		slip = (aW - bW) / totalW
	end
	local dTq = slip * stiffness * c * slipTorque

	return tq * (1 - biasAB) - dTq, tq * biasAB + dTq
end

---@return number, number
local function _HLSDTorqueSplit(tq, aW, bW, aI, bI, biasAB, preload, stiffness, powerRamp, coastRamp, slipTorque)
	if aW < 0 or bW < 0 then
		return tq * (1 - biasAB), tq * biasAB
	end

	local c = tq > 0 and powerRamp or coastRamp
	local tqFactor = math.clamp(math.abs(tq) / slipTorque, -1, 1)
	local totalW = math.abs(aW) + math.abs(bW)
	local slip = 0
	if aW > 0 or bW > 0 then
		slip = (aW - bW) / totalW
	end
	local dTq = slip * stiffness * c * slipTorque * tqFactor

	return tq * (1 - biasAB) - dTq, tq * biasAB + dTq
end

---@enum DIFFERENTIAL_TYPE
DIFFERENTIAL_TYPES = {
	Open = 'Open',
	Locking = 'Locking',
	VLSD = 'VLSD',
	HLSD = 'HLSD',
}

---@type { [DIFFERENTIAL_TYPE]: DifferentialSplitStrategyFn }
local SPLIT_STRATEGIES = {
	[DIFFERENTIAL_TYPES.Open] = _openDiffTorqueSplit,
	[DIFFERENTIAL_TYPES.Locking] = _lockingDiffTorqueSplit,
	[DIFFERENTIAL_TYPES.VLSD] = _VLSDTorqueSplit,
	[DIFFERENTIAL_TYPES.HLSD] = _HLSDTorqueSplit,
}

---@param type DIFFERENTIAL_TYPE
---@return DifferentialSplitStrategyFn
Differential.getSplitStrategy = function(type)
	return SPLIT_STRATEGIES[type] or SPLIT_STRATEGIES[DIFFERENTIAL_TYPES.Open]
end

Differential.TYPES = DIFFERENTIAL_TYPES

---@type fun(vehicle: KPTLVehicle, name?: string, config?: DifferentialConfig): Differential
Differential.new = Differential.new

return Differential