--@include ./friction.txt
--@include ./friction_preset.txt
--@include /koptilnya/libs/constants.txt
--@include /koptilnya/libs/utils.txt

local Friction = require('./friction.txt')
local FrictionPreset = require('./friction_preset.txt')

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

---@class CustomWheelConfig
---@field Name? string
---@field Direction? integer

---@class CustomWheel
local Wheel = class('Wheel')

function Wheel:initialize(config)
	config = config or {}

	self.name = config.Name
	self.direction = config.Direction or 1
	self.mass = config.Mass or 20
	self.radius = config.Radius or 0.27
	self.rollingResistance = config.RollingResistance or 20
	self.squat = config.Squat or 0.1
	self.slipCircleShape = config.SlipCircleShape or 1.05
	self.casterAngle = math.rad(config.CasterAngle or 0)
	self.camberAngle = math.rad(config.CamberAngle or 0)
	self.toeAngle = math.rad(config.ToeAngle or 0)

	self.forwardFriction = Friction:new(config.ForwardFriction)
	self.sideFriction = Friction:new(config.SideFriction)

	self.lateralFrictionPreset = FrictionPreset:new(config.LateralFrictionPreset)
	self.longitudinalFrictionPreset = FrictionPreset:new(config.LongitudinalFrictionPreset)
	self.aligningFrictionPreset = FrictionPreset:new(config.AligningFrictionPreset)

	self.motorTorque = 0
	self.brakeTorque = 0
	self.steerAngle = 0
	self.hasHit = false
	self.counterTorque = 0
	self.inertia = 0
	self.load = 0
	self.angularVelocity = 0
	self.mz = 0

	self.forward = Vector(0)
	self.right = Vector(0)
	self.up = Vector(0)
	self.entity = NULL_ENTITY
	self.physObj = nil
	self.baseInertia = 0.5 * self.mass * math.pow(self.radius, 2)
	self.inertia = self.baseInertia

	self.printDebounced = debounce(function(...)
		print(...)
	end, 1)
end

function Wheel:setEntity(entity)
	self.entity = entity
	self.entity:setNoDraw(false)

	self.physObj = isValid(entity) and entity:getPhysicsObject() or nil

	if isValid(self.physObj) then
		self.physObj:setMaterial('friction_00')
		self.physObj:setMass(self.mass)
		self.physObj:enableDrag(false)
		self.physObj:setDragCoefficient(0)
		self.physObj:setAngleDragCoefficient(0)
	end
end

function Wheel:setInertia(inertia)
	if isValid(self.physObj) then
		print(inertia)
		self.physObj:setInertia(Vector(inertia))
		self.inertia = inertia
	end
end

function Wheel:isValid()
	return isValid(self.entity) and isValid(self.physObj)
end

function Wheel:getRPM()
	return self.angularVelocity * RAD_TO_RPM
end

function Wheel:getLongitudinalLoadCoefficient(load)
	return 11000 * (1 - math.exp(-0.00014 * load))
end

function Wheel:getLateralLoadCoefficient(load)
	return 18000 * (1 - math.exp(-0.0001 * load))
end

function Wheel:stepLongitudinal(Tm, Tb, Vx, W, Lc, R, I)
	local Winit = W
	local VxAbs = math.abs(Vx)
	local Sx = 0

	if Lc / 1000 < 0.01 then
		Sx = 0
	elseif VxAbs >= 0.01 then
		Sx = (W * R - Vx) / VxAbs
	else
		Sx = (W * R - Vx) * 0.6
	end

	Sx = math.clamp(Sx, -1, 1)

	W = W + Tm / I * TICK_INTERVAL

	Tb = Tb * (W > 0 and -1 or 1)
	local TbCap = math.abs(W) * I / TICK_INTERVAL
	local Tbr = math.abs(Tb) - math.abs(TbCap)
	Tbr = math.max(Tbr, 0)
	Tb = math.clamp(Tb, -TbCap, TbCap)
	W = W + Tb / I * TICK_INTERVAL

	local maxTorque = self.longitudinalFrictionPreset:evaluate(math.abs(Sx)) * Lc * R

	local errorTorque = (W - Vx / R) * I / TICK_INTERVAL
	local surfaceTorque = math.clamp(errorTorque, -maxTorque, maxTorque)

	if self.name == 'WheelFL' or self.name == 'WheelFR' then
		surfaceTorque = surfaceTorque
	end

	W = W - surfaceTorque / I * TICK_INTERVAL
	local Fx = surfaceTorque / R

	Tbr = Tbr * (W > 0 and -1 or 1)
	local TbCap2 = math.abs(W) * I / TICK_INTERVAL
	local Tb2 = math.clamp(Tbr, -TbCap2, TbCap2)
	W = W + Tb2 / I * TICK_INTERVAL

	local deltaOmegaTorque = (W - Winit) * I / TICK_INTERVAL
	local Tcnt = -surfaceTorque + Tb + Tb2 - deltaOmegaTorque

	-- self.printDebounced(
	-- 	Color(255, 255, 0),
	-- 	string.format('[%s] ', self.name),
	-- 	Color(255, 255, 255),
	-- 	string.format('%s | %s | %s', math.round(Fx), math.round(W), math.round(R))
	-- )

	return W, Sx, Fx, Tcnt
end

function Wheel:stepLongitudinal2(Tm, Tb, Vx, W, Lc, R, I)
	-- Параметры Pacejka (примерные значения, можно калибровать под задачу)
	local B = 10.0 -- stiffness factor
	local C = 1.9 -- shape factor
	local D = Lc -- peak factor (максимальное сцепление зависит от нагрузки)
	local E = 0.97 -- curvature factor

	-- Угловая скорость колеса и радиус
	local wheelSpeed = W * R

	-- Slip ratio (проверка на деление на 0)
	local slip
	if math.abs(Vx) < 0.1 then
		slip = 0.0
	else
		slip = (wheelSpeed - Vx) / math.abs(Vx)
	end

	-- Формула Pacejka для расчета силы тяги (longitudinal force)
	local Fx = D * math.sin(C * math.atan(B * slip - E * (B * slip - math.atan(B * slip))))

	-- Добавим моторный момент и торможение, сопротивление качению
	local driveForce = Tm / R
	local brakeForce = Tb / R
	local rollingResistanceTorque = self.rollingResistance * R

	local tractionTorque = Fx * R

	-- Суммарный момент, действующий на колесо
	local netTorque = Tm - Tb - rollingResistanceTorque - tractionTorque

	-- Угловое ускорение
	local angularAcceleration = netTorque / I

	-- Обновляем угловую скорость
	local newAngularVelocity = W + angularAcceleration * TICK_INTERVAL

	-- Суммарная продольная сила (на автомобиль)
	local totalLongitudinalForce = Fx + (Tm - Tb - rollingResistanceTorque) / R

	-- print(newAngularVelocity, slip, totalLongitudinalForce, -tractionTorque)

	return newAngularVelocity, slip, totalLongitudinalForce, tractionTorque
end

function Wheel:stepLateral(Vx, Vy, Lc)
	local VxAbs = math.abs(Vx)
	local Sy = 0

	if Lc / 1000 < 0.01 then
		Sy = 0
	elseif VxAbs > 0.1 then
		Sy = math.deg(math.atan(Vy / VxAbs))
		Sy = Sy / 50
	else
		Sy = Vy * (0.003 / TICK_INTERVAL)
	end

	Sy = Sy -- * 0.95
	Sy = math.clamp(Sy, -1, 1)
	local slipSign = Sy < 0 and -1 or 1
	local Fy = -slipSign * self.lateralFrictionPreset:evaluate(math.abs(Sy)) * Lc

	return Sy, Fy
end

function Wheel:slipCircle(Sx, Sy, Fx, Fy, slipCircleShape)
	local SxAbs = math.abs(Sx)

	if SxAbs > 0.01 then
		local SxClamped = math.clamp(Sx, -1, 1)
		local SyClamped = math.clamp(Sy, -1, 1)
		local combinedSlip = Vector2(SxClamped * slipCircleShape, SyClamped)
		local slipDir = combinedSlip:getNormalized()

		local F = math.sqrt(Fx * Fx + Fy * Fy)
		local absSlipDirY = math.abs(slipDir.y)

		Fy = F * absSlipDirY * (Fy < 0 and -1 or 1)
	end

	return Sx, Sy, Fx, Fy
end

function Wheel:selfAligningTorque(Sy, load)
	if math.abs(Sy) < 0.001 or load < 0.001 then
		return 0
	end

	local B = self.aligningFrictionPreset.B
	local C = self.aligningFrictionPreset.C
	local D = self.aligningFrictionPreset.D
	local E = self.aligningFrictionPreset.E

	local entityWheelDown = self.entity:localToWorld(Vector(0, 0, -self.radius))
	local physWheelDown = self.entity:localToWorld(-self.up * self.radius)
	local mechanicalTrail = entityWheelDown:getDistance(physWheelDown) * 10

	local casterEffect = math.tan(self.casterAngle) * self.direction
	local camberTrailEffect = self.camberAngle * 0.005 * self.direction
	-- local toeSyOffset = math.tan(self.toeAngle * self.direction) * 0.1

	local effectiveTrail = mechanicalTrail * self.radius

	local D_scaled = D * load * effectiveTrail

	-- Sy = Sy + toeSyOffset

	local term = B * Sy - E * (B * Sy - math.atan(B * Sy))
	local Mz = D_scaled * math.sin(C * math.atan(term))

	return Mz * self.direction -- / TICK_INTERVAL
end

function Wheel:rotateAxes()
	local ang = self.entity:getAngles()

	-- Base directions
	local forward = ang:getForward() * self.direction
	local up = ang:getUp()
	local right = -ang:getRight() * self.direction

	-- Step 1: Steer + Toe (rotate forward/right around up)
	local steerAngle = -self.steerAngle
	forward = forward:rotateAroundAxis(up, steerAngle)
	right = right:rotateAroundAxis(up, steerAngle)

	-- Step 2: Caster (rotate forward/up around right)
	local casterAngle = -self.casterAngle -- * self.direction
	forward = forward:rotateAroundAxis(right, nil, casterAngle)
	up = up:rotateAroundAxis(right, nil, casterAngle)

	-- Step 3: Camber (rotate up/right around forward)
	local camberAngle = -self.camberAngle * self.direction
	up = up:rotateAroundAxis(forward, nil, camberAngle)
	right = right:rotateAroundAxis(forward, nil, camberAngle)

	-- Normalize final vectors
	self.forward = forward:getNormalized()
	self.right = right:getNormalized() * self.direction
	self.up = up:getNormalized()
end

function Wheel:update()
	if not isValid(self.entity) or not isValid(self.physObj) then
		return
	end

	local externalStress = self.physObj:getStress()
	local velocity = self.physObj:getVelocity() * UNITS_TO_METERS

	self.hasHit = externalStress ~= 0
	self.load = externalStress * KG_TO_KN * 1000

	self.longitudinalLoadCoefficient = self:getLongitudinalLoadCoefficient(self.load)
	self.lateralLoadCoefficient = self:getLateralLoadCoefficient(self.load)

	self:rotateAxes()

	local forwardSpeed = 0
	local sideSpeed = 0

	if self.hasHit then
		forwardSpeed = velocity:dot(self.forward)
		sideSpeed = velocity:dot(self.right)
	end

	local W, Sx, Fx, CounterTq = self:stepLongitudinal(
		self.motorTorque,
		self.brakeTorque + self.rollingResistance,
		forwardSpeed,
		self.angularVelocity,
		self.longitudinalLoadCoefficient,
		self.radius,
		self.inertia
	)

	local Sy, Fy = self:stepLateral(forwardSpeed, sideSpeed, self.lateralLoadCoefficient)

	Sx, Sy, Fx, Fy = self:slipCircle(Sx, Sy, Fx, Fy, 1.05)

	self.mz = self:selfAligningTorque(Sy, self.load)

	self.angularVelocity = W
	self.counterTorque = CounterTq * self.direction
	self.forwardFriction.slip = Sx
	self.forwardFriction.force = Fx
	self.forwardFriction.speed = forwardSpeed
	self.sideFriction.slip = Sy
	self.sideFriction.force = Fy
	self.sideFriction.speed = sideSpeed
end

return Wheel