--!nonstrict
local RunService = game:GetService("RunService")

local avatarJointUpgrade = game:GetEngineFeature("AvatarJointUpgradeFeature")

local Rigging = {}

-- Gravity that joint friction values were tuned under.
local REFERENCE_GRAVITY = 196.2

-- ReferenceMass values from mass of child part. Used to normalized "stiffness" for differently
-- sized avatars (with different mass).
local DEFAULT_MAX_FRICTION_TORQUE = 500
local DEFAULT_STRENGTH_TORQUE = 500

-- Joint friction gets multiplied by this factor during non-graphic death to minimize realistic ragdoll motion
local NONGRAPHIC_FRICTION_SCALE = 100

local HEAD_LIMITS = {
	UpperAngle = 45,
	TwistLowerAngle = -40,
	TwistUpperAngle = 40,
	FrictionTorque = 400,
	StrengthTorque = 400,
	ReferenceMass = 1.0249234437943,
}

local WAIST_LIMITS = {
	UpperAngle = 20,
	TwistLowerAngle = -40,
	TwistUpperAngle = 20,
	FrictionTorque = 750,
	StrengthTorque = 5000, -- Make the waist artificially strong for now. Compensate for the rapidly accelerating RootPart physical model.
	ReferenceMass = 2.861558675766,
}

local ANKLE_LIMITS = {
	UpperAngle = 10,
	TwistLowerAngle = -10,
	TwistUpperAngle = 10,
	StrengthTorque = 400,
	ReferenceMass = 0.43671694397926,
}

local ELBOW_LIMITS = {
	-- Elbow is basically a hinge, but allow some twist for Supination and Pronation
	UpperAngle = 20,
	TwistLowerAngle = 5,
	TwistUpperAngle = 120,
	StrengthTorque = 1200,
	ReferenceMass = 0.70196455717087,
}

local WRIST_LIMITS = {
	UpperAngle = 30,
	TwistLowerAngle = -10,
	TwistUpperAngle = 10,
	StrengthTorque = 400,
	ReferenceMass = 0.69132566452026,
}

local KNEE_LIMITS = {
	UpperAngle = 5,
	TwistLowerAngle = -120,
	TwistUpperAngle = -5,
	StrengthTorque = 2000,
	ReferenceMass = 0.65389388799667,
}

local SHOULDER_LIMITS = {
	UpperAngle = 110,
	TwistLowerAngle = -85,
	TwistUpperAngle = 85,
	FrictionTorque = 600,
	StrengthTorque = 4000,
	ReferenceMass = 0.90918225049973,
}

local HIP_LIMITS = {
	UpperAngle = 40,
	TwistLowerAngle = -5,
	TwistUpperAngle = 80,
	FrictionTorque = 600,
	StrengthTorque = 4000,
	ReferenceMass = 1.9175016880035,
}

local R6_HEAD_LIMITS = {
	UpperAngle = 30,
	TwistLowerAngle = -40,
	TwistUpperAngle = 40,
}

local R6_SHOULDER_LIMITS = {
	UpperAngle = 110,
	TwistLowerAngle = -85,
	TwistUpperAngle = 85,
}

local R6_HIP_LIMITS = {
	UpperAngle = 40,
	TwistLowerAngle = -5,
	TwistUpperAngle = 80,
}

local V3_ZERO = Vector3.new()
local V3_UP = Vector3.new(0, 1, 0)
local V3_DOWN = Vector3.new(0, -1, 0)
local V3_RIGHT = Vector3.new(1, 0, 0)
local V3_LEFT = Vector3.new(-1, 0, 0)

-- To model shoulder cone and twist limits correctly we really need the primary axis of the UpperArm
-- to be going down the limb. the waist and neck joints attachments actually have the same problem
-- of non-ideal axis orientation, but it's not as noticable there since the limits for natural
-- motion are tighter for those joints anyway.
local R15_ADDITIONAL_ATTACHMENTS = {
	{"UpperTorso", "RightShoulderRagdollAttachment", CFrame.fromMatrix(V3_ZERO, V3_RIGHT, V3_UP), "RightShoulderRigAttachment"},
	{"RightUpperArm", "RightShoulderRagdollAttachment", CFrame.fromMatrix(V3_ZERO, V3_DOWN, V3_RIGHT), "RightShoulderRigAttachment"},
	{"UpperTorso", "LeftShoulderRagdollAttachment", CFrame.fromMatrix(V3_ZERO, V3_LEFT, V3_UP), "LeftShoulderRigAttachment"},
	{"LeftUpperArm", "LeftShoulderRagdollAttachment", CFrame.fromMatrix(V3_ZERO, V3_DOWN, V3_LEFT), "LeftShoulderRigAttachment"},
}
-- { { Part0 name (parent), Part1 name (child, parent of joint), attachmentName, limits }, ... }
local R15_RAGDOLL_RIG = {
	{"UpperTorso", "Head", "NeckRigAttachment", HEAD_LIMITS},

	{"LowerTorso", "UpperTorso", "WaistRigAttachment", WAIST_LIMITS},

	{"UpperTorso", "LeftUpperArm", "LeftShoulderRagdollAttachment", SHOULDER_LIMITS},
	{"LeftUpperArm", "LeftLowerArm", "LeftElbowRigAttachment", ELBOW_LIMITS},
	{"LeftLowerArm", "LeftHand", "LeftWristRigAttachment", WRIST_LIMITS},

	{"UpperTorso", "RightUpperArm", "RightShoulderRagdollAttachment", SHOULDER_LIMITS},
	{"RightUpperArm", "RightLowerArm", "RightElbowRigAttachment", ELBOW_LIMITS},
	{"RightLowerArm", "RightHand", "RightWristRigAttachment", WRIST_LIMITS},

	{"LowerTorso", "LeftUpperLeg", "LeftHipRigAttachment", HIP_LIMITS},
	{"LeftUpperLeg", "LeftLowerLeg", "LeftKneeRigAttachment", KNEE_LIMITS},
	{"LeftLowerLeg", "LeftFoot", "LeftAnkleRigAttachment", ANKLE_LIMITS},

	{"LowerTorso", "RightUpperLeg", "RightHipRigAttachment", HIP_LIMITS},
	{"RightUpperLeg", "RightLowerLeg", "RightKneeRigAttachment", KNEE_LIMITS},
	{"RightLowerLeg", "RightFoot", "RightAnkleRigAttachment", ANKLE_LIMITS},
}
-- { { Part0 name, Part1 name }, ... }
local R15_NO_COLLIDES = {
	{"LowerTorso", "LeftUpperArm"},
	{"LeftUpperArm", "LeftHand"},

	{"LowerTorso", "RightUpperArm"},
	{"RightUpperArm", "RightHand"},

	{"LeftUpperLeg", "RightUpperLeg"},

	{"UpperTorso", "RightUpperLeg"},
	{"RightUpperLeg", "RightFoot"},

	{"UpperTorso", "LeftUpperLeg"},
	{"LeftUpperLeg", "LeftFoot"},

	-- Support weird R15 rigs
	{"UpperTorso", "LeftLowerLeg"},
	{"UpperTorso", "RightLowerLeg"},
	{"LowerTorso", "LeftLowerLeg"},
	{"LowerTorso", "RightLowerLeg"},

	{"UpperTorso", "LeftLowerArm"},
	{"UpperTorso", "RightLowerArm"},

	{"Head", "LeftUpperArm"},
	{"Head", "RightUpperArm"},
}
-- { { Motor6D name, Part name }, ...}, must be in tree order, important for ApplyJointVelocities
local R15_MOTOR6DS = {
	{"Waist", "UpperTorso"},

	{"Neck", "Head"},

	{"LeftShoulder", "LeftUpperArm"},
	{"LeftElbow", "LeftLowerArm"},
	{"LeftWrist", "LeftHand"},

	{"RightShoulder", "RightUpperArm"},
	{"RightElbow", "RightLowerArm"},
	{"RightWrist", "RightHand"},

	{"LeftHip", "LeftUpperLeg"},
	{"LeftKnee", "LeftLowerLeg"},
	{"LeftAnkle", "LeftFoot"},

	{"RightHip", "RightUpperLeg"},
	{"RightKnee", "RightLowerLeg"},
	{"RightAnkle", "RightFoot"},
}

-- R6 has hard coded part sizes and does not have a full set of rig Attachments.
local R6_ADDITIONAL_ATTACHMENTS = {
	{"Head", "NeckAttachment", CFrame.new(0, -0.5, 0)},

	{"Torso", "RightShoulderRagdollAttachment", CFrame.fromMatrix(Vector3.new(1, 0.5, 0), V3_RIGHT, V3_UP)},
	{"Right Arm", "RightShoulderRagdollAttachment", CFrame.fromMatrix(Vector3.new(-0.5, 0.5, 0), V3_DOWN, V3_RIGHT)},

	{"Torso", "LeftShoulderRagdollAttachment", CFrame.fromMatrix(Vector3.new(-1, 0.5, 0), V3_LEFT, V3_UP)},
	{"Left Arm", "LeftShoulderRagdollAttachment", CFrame.fromMatrix(Vector3.new(0.5, 0.5, 0), V3_DOWN, V3_LEFT)},

	{"Torso", "RightHipAttachment", CFrame.new(0.5, -1, 0)},
	{"Right Leg", "RightHipAttachment", CFrame.new(0, 1, 0)},

	{"Torso", "LeftHipAttachment", CFrame.new(-0.5, -1, 0)},
	{"Left Leg", "LeftHipAttachment", CFrame.new(0, 1, 0)},
}
-- R6 rig tables use the same table structures as R15.
local R6_RAGDOLL_RIG = {
	{"Torso", "Head", "NeckAttachment", R6_HEAD_LIMITS},

	{"Torso", "Left Leg", "LeftHipAttachment", R6_HIP_LIMITS},
	{"Torso", "Right Leg", "RightHipAttachment", R6_HIP_LIMITS},

	{"Torso", "Left Arm", "LeftShoulderRagdollAttachment", R6_SHOULDER_LIMITS},
	{"Torso", "Right Arm", "RightShoulderRagdollAttachment", R6_SHOULDER_LIMITS},
}
local R6_NO_COLLIDES = {
	{"Left Leg", "Right Leg"},
	{"Head", "Right Arm"},
	{"Head", "Left Arm"},
}
local R6_MOTOR6DS = {
	{"Neck", "Torso"},
	{"Left Shoulder", "Torso"},
	{"Right Shoulder", "Torso"},
	{"Left Hip", "Torso"},
	{"Right Hip", "Torso"},
}

local BALL_SOCKET_NAME = "RagdollBallSocket"
local NO_COLLIDE_NAME = "RagdollNoCollision"

-- Index parts by name to save us from many O(n) FindFirstChild searches
local function indexParts(model)
	local parts = {}
	for _, child in ipairs(model:GetChildren()) do
		if child:IsA("BasePart") then
			local name = child.name
			-- Index first, mimicing FindFirstChild
			if not parts[name] then
				parts[name] = child
			end
		end
	end
	return parts
end

local function createRigJoints(parts, rig, jointUpgradeActive)
	if jointUpgradeActive then
		local rootAC = Instance.new("AnimationConstraint")
		rootAC.Name = "Root"
		rootAC.Attachment0 = parts.HumanoidRootPart.RootRigAttachment
		rootAC.Attachment1 = parts.LowerTorso.RootRigAttachment
		rootAC.IsKinematic = true
		parts.LowerTorso.Root:Destroy()
		rootAC.Parent = parts.LowerTorso
	end

	for _, params in ipairs(rig) do
		local part0Name, part1Name, attachmentName, limits = unpack(params)
		local part0 = parts[part0Name]
		local part1 = parts[part1Name]
		if part0 and part1 then
			local a0 = part0:FindFirstChild(attachmentName)
			local a1 = part1:FindFirstChild(attachmentName)
			if a0 and a1 and a0:IsA("Attachment") and a1:IsA("Attachment") then
				-- Our rigs only have one joint per part (connecting each part to it's parent part), so
				-- we can re-use it if we have to re-rig that part again.
				local constraint = part1:FindFirstChild(BALL_SOCKET_NAME)
				if not constraint then
					constraint = Instance.new("BallSocketConstraint")
					constraint.Name = BALL_SOCKET_NAME
				end
				constraint.Attachment0 = a0
				constraint.Attachment1 = a1
				constraint.LimitsEnabled = true
				constraint.UpperAngle = limits.UpperAngle
				constraint.TwistLimitsEnabled = true
				constraint.TwistLowerAngle = limits.TwistLowerAngle
				constraint.TwistUpperAngle = limits.TwistUpperAngle
				-- Scale constant torque limit for joint friction relative to gravity and the mass of
				-- the body part.
				local gravityScale = workspace.Gravity / REFERENCE_GRAVITY
				if jointUpgradeActive then
					gravityScale = math.max(gravityScale, .2) -- Still need some joint friction in zero G
				end
				local referenceMass = limits.ReferenceMass
				local massScale = referenceMass and (part1:GetMass() / referenceMass) or 1
				local maxTorque = limits.FrictionTorque or DEFAULT_MAX_FRICTION_TORQUE
				constraint.MaxFrictionTorque = maxTorque * massScale * gravityScale
				if jointUpgradeActive then
					constraint.MaxFrictionTorque /= NONGRAPHIC_FRICTION_SCALE
				end
				constraint.Parent = part1

				if jointUpgradeActive then
					local motor = part1:FindFirstChildWhichIsA("Motor6D")
					assert(motor)

					local ac = Instance.new("AnimationConstraint")
					ac.Name = motor.Name
					ac.Attachment0 = motor.Part0:FindFirstChild(motor.Name .. "RigAttachment")
					ac.Attachment1 = motor.Part1:FindFirstChild(motor.Name .. "RigAttachment")
					ac.IsKinematic = true
					local strengthTorque = limits.StrengthTorque or DEFAULT_STRENGTH_TORQUE
					ac.MaxTorque = strengthTorque * massScale * gravityScale
					ac.MaxForce = 0 -- currently simulated joints don't allow translation
					ac.Parent = motor.parent
					motor:Destroy()
				end
			end
		end
	end
end

local function createAdditionalAttachments(parts, attachments)
	for _, attachmentParams in ipairs(attachments) do
		local partName, attachmentName, cframe, baseAttachmentName = unpack(attachmentParams)
		local part = parts[partName]
		if part then
			local attachment = part:FindFirstChild(attachmentName)
			-- Create or update existing attachment
			if not attachment or attachment:IsA("Attachment") then
				if baseAttachmentName then
					local base = part:FindFirstChild(baseAttachmentName)
					if base and base:IsA("Attachment") then
						cframe = base.CFrame * cframe
					end
				end
				-- The attachment names are unique within a part, so we can re-use
				if not attachment then
					attachment = Instance.new("Attachment")
					attachment.Name = attachmentName
					attachment.CFrame = cframe
					attachment.Parent = part
				else
					attachment.CFrame = cframe
				end
			end
		end
	end
end

local function createNoCollides(parts, noCollides)
	-- This one's trickier to handle for an already rigged character since a part will have multiple
	-- NoCollide children with the same name. Having fewer unique names is better for
	-- replication so we suck it up and deal with the complexity here.

	-- { [Part1] = { [Part0] = true, ... }, ...}
	local needed = {}
	-- Following the convention of the Motor6Ds and everything else here we parent the NoCollide to
	-- Part1, so we start by building the set of Part0s we need a NoCollide with for each Part1
	for _, namePair in ipairs(noCollides) do
		local part0Name, part1Name = unpack(namePair)
		local p0, p1 = parts[part0Name], parts[part1Name]
		if p0 and p1 then
			local p0Set = needed[p1]
			if not p0Set then
				p0Set = {}
				needed[p1] = p0Set
			end
			p0Set[p0] = true
		end
	end

	-- Go through NoCollides that exist and remove Part0s from the needed set if we already have
	-- them covered. Gather NoCollides that aren't between parts in the set for resue
	local reusableNoCollides = {}
	for part1, neededPart0s in pairs(needed) do
		local reusables = {}
		for _, child in ipairs(part1:GetChildren()) do
			if child:IsA("NoCollisionConstraint") and child.Name == NO_COLLIDE_NAME then
				local p0 = child.Part0
				local p1 = child.Part1
				if p1 == part1 and neededPart0s[p0] then
					-- If this matches one that we needed, we don't need to create it anymore.
					neededPart0s[p0] = nil
				else
					-- Otherwise we're free to reuse this NoCollide
					table.insert(reusables, child)
				end
			end
		end
		reusableNoCollides[part1] = reusables
	end

	-- Create the remaining NoCollides needed, re-using old ones if possible
	for part1, neededPart0s in pairs(needed) do
		local reusables = reusableNoCollides[part1]
		for part0, _ in pairs(neededPart0s) do
			local constraint = table.remove(reusables)
			if not constraint then
				constraint = Instance.new("NoCollisionConstraint")
			end
			constraint.Name = NO_COLLIDE_NAME
			constraint.Part0 = part0
			constraint.Part1 = part1
			constraint.Parent = part1
		end
	end
end

local function isAttachmetInPart(attachment, part)
	return attachment and attachment.Parent == part or false
end

local function checkValidAttachment(part0, part1, attachment0, attachment1)
	if isAttachmetInPart(attachment0, part0) and isAttachmetInPart(attachment1, part1) then
		return true
	end

	if isAttachmetInPart(attachment0, part1) and isAttachmetInPart(attachment1, part0) then
		return true
	end

	return false
end

local function hasValidConstraint(part0, part1)
	for _, child in ipairs(part1:GetChildren()) do
		if child:IsA("BallSocketConstraint") then
			local attachment0 = child.Attachment0
			local attachment1 = child.Attachment1

			if checkValidAttachment(part0, part1, attachment0, attachment1) then
				return true
			end
		end
	end

	return false
end

local function hasRagdollJoint(motor : Instance)
	if avatarJointUpgrade and not (motor:IsA("Motor6D") or motor:IsA("AnimationConstraint")) then
		return false
	end
	local part0
	local part1
	if motor:IsA("Motor6D") then
		part0 = motor.Part0
		part1 = motor.Part1
	elseif avatarJointUpgrade and motor:IsA("AnimationConstraint") then
		local at0 = motor.Attachment0
		local at1 = motor.Attachment1
		part0 = at0 and at0.Parent
		part1 = at1 and at1.Parent
	end

	if not part0 or not part1 then
		return false
	end

	if hasValidConstraint(part0, part1) then
		return true
	end

	-- Don't enforce ordering for developer created ragdolls
	if hasValidConstraint(part1, part0) then
		return true
	end

	return false
end

local function disableMotorSet(model, motorSet)
	local motors = {}
	-- Destroy all regular joints:
	for _, params in ipairs(motorSet) do
		local part = model:FindFirstChild(params[2])
		if part then
			local motor = part:FindFirstChild(params[1])
			if motor and (avatarJointUpgrade or motor:IsA("Motor6D")) and hasRagdollJoint(motor) then
				table.insert(motors, motor)
				motor.Enabled = false
			end
		end
	end
	return motors
end

function Rigging.createRagdollJoints(model, rigType, jointUpgradeActive)
	local parts = indexParts(model)
	if rigType == Enum.HumanoidRigType.R6 then
		createAdditionalAttachments(parts, R6_ADDITIONAL_ATTACHMENTS)
		createRigJoints(parts, R6_RAGDOLL_RIG, false)
		createNoCollides(parts, R6_NO_COLLIDES)
	elseif rigType == Enum.HumanoidRigType.R15 then
		createAdditionalAttachments(parts, R15_ADDITIONAL_ATTACHMENTS)
		createRigJoints(parts, R15_RAGDOLL_RIG, jointUpgradeActive)
		createNoCollides(parts, R15_NO_COLLIDES)
	else
		error("unknown rig type", 2)
	end
end

function Rigging.disableMotors(model, rigType)
	-- Note: We intentionally do not disable the root joint so that the mechanism root of the
	-- character stays consistent when we break joints on the client. This avoid the need for the client to wait
	-- for re-assignment of network ownership of a new mechanism, which creates a visible hitch.

	local motors
	if rigType == Enum.HumanoidRigType.R6 then
		motors = disableMotorSet(model, R6_MOTOR6DS)
	elseif rigType == Enum.HumanoidRigType.R15 then
		motors = disableMotorSet(model, R15_MOTOR6DS)
	else
		error("unknown rig type", 2)
	end

	-- Set the root part to non-collide
	local rootPart = model.PrimaryPart or model:FindFirstChild("HumanoidRootPart")
	if rootPart and rootPart:IsA("BasePart") then
		rootPart.CanCollide = false
	end

	-- Set the Head part to collide
	local head = model:FindFirstChild("Head")
	if head and head:IsA("BasePart") then
		head.CanCollide = true
	end

	return motors
end

function Rigging.disableParticleEmittersAndFadeOut(character, duration)
	if RunService:IsServer() then
		-- This causes a lot of unnecesarry replicated property changes
		error("disableParticleEmittersAndFadeOut should not be called on the server.", 2)
	end

	local descendants = character:GetDescendants()
	local transparencies = {}
	for _, instance in pairs(descendants) do
		if instance:IsA("BasePart") or instance:IsA("Decal") then
			transparencies[instance] = instance.Transparency
		elseif instance:IsA("ParticleEmitter") then
			instance.Enabled = false
		end
	end
	local t = 0
	while t < duration do
		-- Using heartbeat because we want to update just before rendering next frame, and not
		-- block the render thread kicking off (as RenderStepped does)
		local dt = RunService.Heartbeat:Wait()
		t = t + dt
		local alpha = math.min(t / duration, 1)
		for part, initialTransparency in pairs(transparencies) do
			part.Transparency = (1 - alpha) * initialTransparency + alpha
		end
	end
end

function Rigging.easeNongraphicJointFriction(character, duration)
	local descendants = character:GetDescendants()
	-- { { joint, initial friction, end friction }, ... }
	local frictionJoints = {}
	for _, v in pairs(descendants) do
		if v:IsA("BallSocketConstraint") and v.Name == BALL_SOCKET_NAME then
			v.MaxFrictionTorque *= NONGRAPHIC_FRICTION_SCALE
			local current = v.MaxFrictionTorque
			-- Keep the torso and neck a little stiffer...
			local parentName = v.Parent.Name
			local scale = (parentName == "UpperTorso" or parentName == "Head") and 0.5 or 0.05
			local nextTorque = current * scale
			frictionJoints[v] = { v, current, nextTorque }
		end
	end
	local t = 0
	while t < duration do
		-- Using stepped because we want to update just before physics sim
		local _, dt = RunService.Stepped:Wait()
		t = t + dt
		local alpha = math.min(t / duration, 1)
		for _, tuple in pairs(frictionJoints) do
			local ballSocket, a, b = unpack(tuple)
			ballSocket.MaxFrictionTorque = (1 - alpha) * a + alpha * b
		end
	end
end

-- remove with FFlag::AvatarJointUpgrade
function Rigging.easeJointFriction_OLD(character, duration)
	local descendants = character:GetDescendants()
	-- { { joint, initial friction, end friction }, ... }
	local frictionJoints = {}
	for _, v in pairs(descendants) do
		if v:IsA("BallSocketConstraint") and v.Name == BALL_SOCKET_NAME then
			local current = v.MaxFrictionTorque
			-- Keep the torso and neck a little stiffer...
			local parentName = v.Parent.Name
			local scale = (parentName == "UpperTorso" or parentName == "Head") and 0.5 or 0.05
			local nextTorque = current * scale
			frictionJoints[v] = { v, current, nextTorque }
		end
	end
	local t = 0
	while t < duration do
		-- Using stepped because we want to update just before physics sim
		local _, dt = RunService.Stepped:Wait()
		t = t + dt
		local alpha = math.min(t / duration, 1)
		for _, tuple in pairs(frictionJoints) do
			local ballSocket, a, b = unpack(tuple)
			ballSocket.MaxFrictionTorque = (1 - alpha) * a + alpha * b
		end
	end
end

return Rigging
