/* Copyright (C) Itseez3D, Inc. - All Rights Reserved
* You may not use this file except in compliance with an authorized license
* Unauthorized copying of this file, via any medium is strictly prohibited
* Proprietary and confidential
* UNLESS REQUIRED BY APPLICABLE LAW OR AGREED BY ITSEEZ3D, INC. IN WRITING, SOFTWARE DISTRIBUTED UNDER THE LICENSE IS DISTRIBUTED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OR
* CONDITIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED
* See the License for the specific language governing permissions and limitations under the License.
* Written by Itseez3D, Inc. <support@avatarsdk.com>, June 2026
*/

using System.Collections.Generic;
using System.Reflection;
using System;
using UnityEngine;

namespace AvatarSDKMove
{
	public sealed class HumanDescriptionBoneCache
	{
		public readonly struct BoneRotationInfo
		{
			public readonly HumanBodyBones Bone;
			public readonly Transform Transform;

			// World rotation reconstructed from HumanDescription.skeleton (T-pose)
			public readonly Quaternion SkeletonWorldRotation;

			// Parent world rotation reconstructed from HumanDescription.skeleton (T-pose)
			public readonly Quaternion ParentSkeletonWorldRotation;

			public BoneRotationInfo(
				HumanBodyBones bone,
				Transform transform,
				Quaternion skeletonWorldRotation,
				Quaternion parentSkeletonWorldRotation)
			{
				Bone = bone;
				Transform = transform;
				SkeletonWorldRotation = skeletonWorldRotation;
				ParentSkeletonWorldRotation = parentSkeletonWorldRotation;
			}
		}

		private readonly Dictionary<HumanBodyBones, BoneRotationInfo> _bones = new();

		public IReadOnlyDictionary<HumanBodyBones, BoneRotationInfo> Bones => _bones;

		public HumanDescriptionBoneCache(Transform root, HumanDescription humanDescription)
		{
			if (root == null) throw new ArgumentNullException(nameof(root));

			Build(root, humanDescription);
		}

		public bool TryGet(HumanBodyBones bone, out BoneRotationInfo info)
		{
			return _bones.TryGetValue(bone, out info);
		}

		private void Build(Transform root, HumanDescription humanDescription)
		{
			_bones.Clear();

			if (humanDescription.human == null || humanDescription.skeleton == null)
				return;

			// Human bone name -> actual transform name/path
			var humanToBoneName = new Dictionary<string, string>(StringComparer.Ordinal);
			foreach (var humanBone in humanDescription.human)
			{
				if (!string.IsNullOrEmpty(humanBone.humanName) && !string.IsNullOrEmpty(humanBone.boneName))
					humanToBoneName[humanBone.humanName] = humanBone.boneName;
			}

			// Skeleton bone name -> SkeletonBone
			var skeletonByName = new Dictionary<string, SkeletonBone>(StringComparer.Ordinal);
			foreach (var skeletonBone in humanDescription.skeleton)
			{
				if (!string.IsNullOrEmpty(skeletonBone.name))
					skeletonByName[skeletonBone.name] = skeletonBone;
			}

			// Cache world rotations computed from the skeleton T-pose hierarchy.
			var worldRotationCache = new Dictionary<string, Quaternion>(StringComparer.Ordinal);

			Quaternion GetSkeletonWorldRotation(string boneName)
			{
				if (string.IsNullOrEmpty(boneName))
					return Quaternion.identity;

				if (worldRotationCache.TryGetValue(boneName, out var cached))
					return cached;

				if (!skeletonByName.TryGetValue(boneName, out var skeletonBone))
				{
					worldRotationCache[boneName] = Quaternion.identity;
					return Quaternion.identity;
				}

				string parentName = GetSkeletonParentName(skeletonBone);

				Quaternion localRotation = skeletonBone.rotation;
				Quaternion worldRotation = localRotation;

				if (!string.IsNullOrEmpty(parentName) && skeletonByName.ContainsKey(parentName))
				{
					Quaternion parentWorldRotation = GetSkeletonWorldRotation(parentName);
					worldRotation = parentWorldRotation * localRotation;
				}

				worldRotationCache[boneName] = worldRotation;
				return worldRotation;
			}

			Quaternion GetSkeletonParentWorldRotation(string boneName)
			{
				if (string.IsNullOrEmpty(boneName))
					return Quaternion.identity;

				if (!skeletonByName.TryGetValue(boneName, out var skeletonBone))
					return Quaternion.identity;

				string parentName = GetSkeletonParentName(skeletonBone);
				if (string.IsNullOrEmpty(parentName) || !skeletonByName.ContainsKey(parentName))
					return Quaternion.identity;

				return GetSkeletonWorldRotation(parentName);
			}

			foreach (HumanBodyBones bone in Enum.GetValues(typeof(HumanBodyBones)))
			{
				if (bone == HumanBodyBones.LastBone)
					continue;

				string humanName = bone.ToString();
				if (!humanToBoneName.TryGetValue(humanName, out string boneName))
					continue;

				Transform tr = FindTransform(root, boneName);
				if (tr == null)
					continue;

				Quaternion worldRotation = GetSkeletonWorldRotation(boneName);
				Quaternion parentWorldRotation = GetSkeletonParentWorldRotation(boneName);

				_bones[bone] = new BoneRotationInfo(
					bone,
					tr,
					worldRotation,
					parentWorldRotation
				);
			}
		}

		private static Transform FindTransform(Transform root, string boneNameOrPath)
		{
			// First try path lookup.
			Transform found = root.Find(boneNameOrPath);
			if (found != null)
				return found;

			// Fallback to recursive name search.
			return FindByNameRecursive(root, boneNameOrPath);
		}

		private static Transform FindByNameRecursive(Transform current, string name)
		{
			if (current.name == name)
				return current;

			for (int i = 0; i < current.childCount; i++)
			{
				Transform child = current.GetChild(i);
				Transform result = FindByNameRecursive(child, name);
				if (result != null)
					return result;
			}

			return null;
		}

		private static readonly FieldInfo SkeletonParentNameField =
			typeof(SkeletonBone).GetField("parentName", BindingFlags.Instance | BindingFlags.NonPublic);

		private static string GetSkeletonParentName(SkeletonBone bone)
		{
			if (SkeletonParentNameField == null)
				return null;

			return SkeletonParentNameField.GetValue(bone) as string;
		}
	}
}
