﻿using System;
using System.Collections.Concurrent;
using System.Collections.Generic;
using System.Threading;
using RebocapWsSdk;
using UnityEngine;
using Random = UnityEngine.Random;

namespace RebocapSdk.DemoScenes
{
    public class SdkManager
    {
        private RebocapWsSdk.RebocapWsSdk _sdk;
        private uint _portNumber;
        private readonly ConcurrentQueue<PoseMsg> _dataQ = new();
        // sync signals for ui or thread
        private bool _gameExist;
        private bool _connectButtonClicked;
        private bool _triggerRegisterAnimator;
        private bool _updattingAvatar;
        
        // skeleton data
        private List<Vector3> _leftFootVertex, _rightFootVertex, _leftFootNorm, _rightFootNorm, _skeletonPosition;

        private Action<bool> _uiUpdate;
        
        private Animator _avatar;
        private bool _connected;
        private bool _isAvatarNull;
        private int _uid;
        private Quaternion _defaultRootRotation;
        private Quaternion[] _defaultLocalRotation;
        private Quaternion[] _defaultBindRotation;
        private Quaternion[] _defaultBindRotationInverse;
        private Transform[] _bones;
        private Transform _rootTransform;
        private bool[] _emptyBones;
        private Vector3 _root2Hip;
        public static int[] RebocapBoneParents =
        {
            -1, // Hip
            (int)RebocapBones.Hip, // LeftUpperLeg,
            (int)RebocapBones.Hip, // RightUpperLeg,
            (int)RebocapBones.Hip, // Spine,
            (int)RebocapBones.LeftUpperLeg, // LeftLowerLeg,
            (int)RebocapBones.RightUpperLeg, // RightLowerLeg,
            (int)RebocapBones.Spine, // Chest,
            (int)RebocapBones.LeftLowerLeg, // LeftFoot,
            (int)RebocapBones.RightLowerLeg, // RightFoot,
            (int)RebocapBones.Chest, // UpperChest,
            (int)RebocapBones.LeftFoot, // LeftToe,
            (int)RebocapBones.RightFoot, // RightToe,
            (int)RebocapBones.UpperChest, // Neck,
            (int)RebocapBones.UpperChest, // LeftShoulder,
            (int)RebocapBones.UpperChest, // RightShoulder,
            (int)RebocapBones.Neck, // Head,
            (int)RebocapBones.LeftShoulder, // LeftUpperArm,
            (int)RebocapBones.RightShoulder, // RightUpperArm,
            (int)RebocapBones.LeftUpperArm, // LeftLowerArm,
            (int)RebocapBones.RightUpperArm, // RightLowerArm,
            (int)RebocapBones.LeftLowerArm, // LeftHand,
            (int)RebocapBones.RightLowerArm, // RightHand,
            (int)RebocapBones.LeftHand, // LeftMiddleProximal,
            (int)RebocapBones.RightHand // RightMiddleProximal
        };

        public SdkManager(uint portNumber, Action<bool> updateUi)
        {
            _portNumber = portNumber;
            _bones = new Transform[24];
            _defaultLocalRotation = new Quaternion[24];
            _defaultBindRotation = new Quaternion[24];
            _defaultBindRotationInverse = new Quaternion[24];
            _emptyBones = new bool[24];
            for (int i = 0; i < 24; i++) _emptyBones[i] = true;
            _connected = false;
            _isAvatarNull = true;
            _triggerRegisterAnimator = false;
            _updattingAvatar = false;
            
            _root2Hip = Vector3.zero;
            _uiUpdate = updateUi;
            _uid = Random.Range(0, 65536);
            
            new Thread(SdkThreadLoop).Start();
        }

        public void TriggerConnectStatusChange()
        {
            _connectButtonClicked = true;
        }

        public void DestroySdk()
        {
            _gameExist = true;
        }

        // all sdk handle should on child thread!!!!
        private void SdkThreadLoop()
        {
            _sdk = new RebocapWsSdk.RebocapWsSdk();
            _sdk.SetPoseMsgCallback(OnDataReceive);
            _sdk.SetExceptionCloseCallback(OnWsException);

            while (!_gameExist)
            {
                if (_connectButtonClicked)
                {
                    _connectButtonClicked = false;
                    if (_connected)
                    {
                        Disconnect();
                    }
                    else
                    {
                        Connect(_portNumber);
                    }
                    ThreadUIHelper.Instance.RunOnMainThread(parameters => _uiUpdate((bool)parameters[0]), _connected);
                }

                if (_triggerRegisterAnimator)
                {
                    _triggerRegisterAnimator = false;
                    RegisterSkeletonToRebocap();
                }
                
                Thread.Sleep(16);
            }
            Disconnect();
        }

        public bool Connect(uint port)
        {
            if (_connected) return false;
            
            if (port < 100 || port > 65535) return false;
            var openRes = _sdk.Open(port, uid:_uid);
            if (openRes == 0)
            {
                _connected = true;
                return true;
            }
            
            Debug.Log("connection failed" + openRes);
            switch (openRes)
            {
                case 1:
                    Debug.Log("connection status error");
                    break;
                case 2:
                    Debug.Log("connection failed");
                    break;
                case 3:
                    Debug.Log("auth failed");
                    break;
                default:
                    Debug.Log("unknown connection error");
                    break;
            }

            return false;
        }

        public void Disconnect()
        {
            if (!_connected) return;
            _connected = false;
            _sdk.Close();
        }

        private void OnWsException(RebocapWsSdk.RebocapWsSdk self)
        {
            Debug.Log("ws exception happened!!! should handle disconnect here");
            _connected = false;
            _uiUpdate(_connected);
        }
        
        private void OnDataReceive(PoseMsg msg, RebocapWsSdk.RebocapWsSdk self)
        {
            if (_updattingAvatar) return;
            if (_dataQ.Count > 3)
            {
                while (_dataQ.Count > 1) _dataQ.TryDequeue(out var _);
            }
            _dataQ.Enqueue(msg);
        }

        // you must connect to register avatar and you should make sure avatar is in T-pose before call this!!!
        public void RegisterAvatar(Animator animator, bool registerToRebocap)
        {
            _updattingAvatar = true;
            while (!_dataQ.IsEmpty)
            {
                _dataQ.TryDequeue(out var msg);
            }
            _avatar = animator;
            _isAvatarNull = false;  // avoid _avatar null detect which cost a lot
            SetBones();
            if (_bones[(int)RebocapBones.Hip] == null)
            {
                Debug.Log("hip bone is root for rebocap capture!!! miss hip bone!!!");
                return;
            }

            if (!registerToRebocap) return;
            // skip Update Pose will set to false after ResetAvatar called, the following code is use to
            // make sure avatar in T-pose to calculate skeleton and register to rebocap
            ResetAvatar();
            // 1. find vertex
            GetFootVertexInfo(out _leftFootVertex, out _rightFootVertex, out _leftFootNorm, out _rightFootNorm);
            
            // calculate T-pose bone position
            _skeletonPosition = new List<Vector3> { _bones[0].position };
            for (var i = 1; i < 24; i++)
            {
                Vector3 p;
                var cBoneID = i;
                // if bone miss, use parent position
                while (_bones[cBoneID] == null)
                {
                    cBoneID = RebocapBoneParents[cBoneID];
                }

                p = _bones[cBoneID].position;
                
                if (i != cBoneID) p += new Vector3(1.0e-6f, 1.0e-6f, 1.0e-6f);
                _skeletonPosition.Add(p);
            }
            _triggerRegisterAnimator = true;
            _updattingAvatar = false;
        }

        private bool RegisterSkeletonToRebocap()
        {
            if (_avatar == null || !_connected) return false; 
            // register skeleton to rebocap
            List<Vector3> bindFootVertices = new List<Vector3>();

#if ENABLE_IL2CPP
            int status = _sdk.CalculateFootVertexAndRegisterAvatarIl2CPP(_skeletonPosition, _leftFootVertex, _rightFootVertex,
                    _leftFootNorm, _rightFootNorm, bindFootVertices, "-xyz");
#else
            int status = _sdk.CalculateFootVertexAndRegisterAvatar(_skeletonPosition, _leftFootVertex, _rightFootVertex,
                    _leftFootNorm, _rightFootNorm, bindFootVertices, "-xyz");
#endif
            if (status == 0)
            {
                return true;
            }
            if (status == -1)
            {
                Debug.Log("network register failed");
            }
            else
            {
                Debug.Log($"other error when register skeleton to rebocap:{status}");
            }
            return false;
        }

        private static void FindMeshVertexByBoneIds(in SkinnedMeshRenderer meshRenderer,
            in List<int> leftIds, in List<int> rightIds,
            List<Vector3> selectedLeftVertices, List<Vector3> selectedLeftNormals,
            List<Vector3> selectedRightVertices, List<Vector3> selectedRightNormals)
        {
            if (leftIds.Count == 0 && rightIds.Count == 0) return;
            
            Mesh mesh = new Mesh();
            meshRenderer.BakeMesh(mesh, true);
            var vertices = new List<Vector3>();
            var normals = new List<Vector3>();
            mesh.GetVertices(vertices);
            mesh.GetNormals(normals);
            var boneWeights = meshRenderer.sharedMesh.boneWeights;
            for (var i = 0; i < vertices.Count; i++)
            {
                vertices[i] = meshRenderer.transform.TransformPoint(vertices[i]);
                if (leftIds.Contains(boneWeights[i].boneIndex0) || leftIds.Contains(boneWeights[i].boneIndex1))
                {
                    selectedLeftVertices.Add(vertices[i]);
                    selectedLeftNormals.Add(normals[i]);
                }
                if (rightIds.Contains(boneWeights[i].boneIndex0) || rightIds.Contains(boneWeights[i].boneIndex1))
                {
                    selectedRightVertices.Add(vertices[i]);
                    selectedRightNormals.Add(normals[i]);
                }
            }
        }
    
        private void GetFootVertexInfo(out List<Vector3> leftFootVertex, out List<Vector3> rightFootVertex,
            out List<Vector3> leftFootNorm, out List<Vector3> rightFootNorm)
        {
            leftFootVertex = new List<Vector3>();
            rightFootVertex = new List<Vector3>();
            leftFootNorm = new List<Vector3>();
            rightFootNorm = new List<Vector3>();
            if (!_isAvatarNull) return;

            var meshRenderers = _avatar.transform.GetComponentsInChildren<SkinnedMeshRenderer>();

            foreach (var meshRenderer in meshRenderers)
            {
                List<int> leftFootIds = new List<int>();
                List<int> rightFootIds = new List<int>();
                for (var i = 0; i < meshRenderer.bones.Length; i++)
                {
                    if (meshRenderer.bones[i].transform.Equals(_bones[(int)RebocapBones.LeftFoot]) ||
                        meshRenderer.bones[i].transform.Equals(_bones[(int)RebocapBones.LeftToe]))
                    {
                        leftFootIds.Add(i);
                    } else if (meshRenderer.bones[i].transform.Equals(_bones[(int)RebocapBones.RightFoot]) ||
                               meshRenderer.bones[i].transform.Equals(_bones[(int)RebocapBones.RightToe]))
                    {
                        rightFootIds.Add(i);
                    }
                }
                FindMeshVertexByBoneIds(meshRenderer, leftFootIds, rightFootIds,
                    leftFootVertex, leftFootNorm,
                    rightFootVertex, rightFootNorm);
            }
            if (leftFootVertex.Count == 0 || rightFootVertex.Count == 0)
            {
                // if foot vertex is still empty, take lower leg as foot
                foreach (var meshRenderer in meshRenderers)
                {
                    List<int> leftFootIds = new List<int>();
                    List<int> rightFootIds = new List<int>();
                    for (var i = 0; i < meshRenderer.bones.Length; i++)
                    {
                        if (meshRenderer.bones[i].transform.Equals(_bones[(int)RebocapBones.LeftLowerLeg]))
                        {
                            leftFootIds.Add(i);
                        } else if (meshRenderer.bones[i].transform.Equals(_bones[(int)RebocapBones.RightLowerLeg]))
                        {
                            rightFootIds.Add(i);
                        }
                    }
                    FindMeshVertexByBoneIds(meshRenderer, leftFootIds, rightFootIds,
                        leftFootVertex, leftFootNorm,
                        rightFootVertex, rightFootNorm);
                }
            }
        }

        private void SetBones()
        {
            if (_isAvatarNull) return;
            // update skeleton 24
            _bones[(uint)RebocapBones.Hip] = _avatar.GetBoneTransform(HumanBodyBones.Hips);
            _bones[(uint)RebocapBones.Spine] = _avatar.GetBoneTransform(HumanBodyBones.Spine);
            _bones[(uint)RebocapBones.Chest] = _avatar.GetBoneTransform(HumanBodyBones.Chest);
            _bones[(uint)RebocapBones.UpperChest] = _avatar.GetBoneTransform(HumanBodyBones.UpperChest);
            _bones[(uint)RebocapBones.Neck] = _avatar.GetBoneTransform(HumanBodyBones.Neck);
            _bones[(uint)RebocapBones.Head] = _avatar.GetBoneTransform(HumanBodyBones.Head);
            _bones[(uint)RebocapBones.LeftShoulder] = _avatar.GetBoneTransform(HumanBodyBones.LeftShoulder);
            _bones[(uint)RebocapBones.RightShoulder] = _avatar.GetBoneTransform(HumanBodyBones.RightShoulder);
            _bones[(uint)RebocapBones.LeftUpperArm] = _avatar.GetBoneTransform(HumanBodyBones.LeftUpperArm);
            _bones[(uint)RebocapBones.RightUpperArm] = _avatar.GetBoneTransform(HumanBodyBones.RightUpperArm);
            _bones[(uint)RebocapBones.LeftLowerArm] = _avatar.GetBoneTransform(HumanBodyBones.LeftLowerArm);
            _bones[(uint)RebocapBones.RightLowerArm] = _avatar.GetBoneTransform(HumanBodyBones.RightLowerArm);
            _bones[(uint)RebocapBones.LeftHand] = _avatar.GetBoneTransform(HumanBodyBones.LeftHand);
            _bones[(uint)RebocapBones.RightHand] = _avatar.GetBoneTransform(HumanBodyBones.RightHand);
            _bones[(uint)RebocapBones.LeftMiddleProximal] = _avatar.GetBoneTransform(HumanBodyBones.LeftMiddleProximal);
            _bones[(uint)RebocapBones.RightMiddleProximal] = _avatar.GetBoneTransform(HumanBodyBones.RightMiddleProximal);
            _bones[(uint)RebocapBones.LeftUpperLeg] = _avatar.GetBoneTransform(HumanBodyBones.LeftUpperLeg);
            _bones[(uint)RebocapBones.RightUpperLeg] = _avatar.GetBoneTransform(HumanBodyBones.RightUpperLeg);
            _bones[(uint)RebocapBones.LeftLowerLeg] = _avatar.GetBoneTransform(HumanBodyBones.LeftLowerLeg);
            _bones[(uint)RebocapBones.RightLowerLeg] = _avatar.GetBoneTransform(HumanBodyBones.RightLowerLeg);
            _bones[(uint)RebocapBones.LeftFoot] = _avatar.GetBoneTransform(HumanBodyBones.LeftFoot);
            _bones[(uint)RebocapBones.RightFoot] = _avatar.GetBoneTransform(HumanBodyBones.RightFoot);
            _bones[(uint)RebocapBones.LeftToe] = _avatar.GetBoneTransform(HumanBodyBones.LeftToes);
            _bones[(uint)RebocapBones.RightToe] = _avatar.GetBoneTransform(HumanBodyBones.RightToes);

            Quaternion armatureRotation = _bones[0].parent.rotation;
            for (int i = 0; i < _bones.Length; i++)
            {
                if (_bones[i] == null)
                {
                    _emptyBones[i] = true;
                }
                else
                {
                    _emptyBones[i] = false;
                    _defaultLocalRotation[i] = _bones[i].localRotation;

                    var parentGlobalRotation = _bones[i].rotation;
                    _defaultBindRotation[i] = _bones[i].rotation;
                    _defaultBindRotationInverse[i] = Quaternion.Inverse(_defaultBindRotation[i]);
                }
            }

            _defaultRootRotation = _avatar.avatarRoot.localRotation;
            if (!_emptyBones[(int)RebocapBones.Hip]) _root2Hip = _bones[(int)RebocapBones.Hip].position - _avatar.avatarRoot.position;
        }

        private void ResetAvatar()
        {
            // make sure avatar is default rotation and set to T-pose
            // if avatar is bind to some other object as child, set it's parent rotation to default!
            if (_isAvatarNull)
            {
                return;
            }
            
            _avatar.avatarRoot.localRotation = _defaultRootRotation;
            for (var i = 0; i < _bones.Length; i++)
            {
                if (_emptyBones[i]) continue;
                _bones[i].localRotation = _defaultLocalRotation[i];
            }
        }

        public void HandleAvatarUpdate()
        {
            if (_isAvatarNull || _dataQ.Count == 0)
            {
                return;
            }

            _dataQ.TryDequeue(out var msg);
            while (_dataQ.TryDequeue(out var tempMsg))
            {
                msg = tempMsg;
            }

            var currentHipPosition = new Vector3(msg.Trans[0], msg.Trans[1], msg.Trans[2]);
            _avatar.avatarRoot.position = currentHipPosition;
            for (var i = 0; i < _bones.Length; i++)
            {
                if (_emptyBones[i]) continue;
                // _bones[i].rotation = _defaultLocalRotation[i] * (_defaultBindRotationInverse[i] * (new Quaternion(msg.Pose24[i, 0], msg.Pose24[i, 1], msg.Pose24[i, 2], msg.Pose24[i, 3])) * _defaultBindRotation[i]);
                _bones[i].rotation = (new Quaternion(msg.Pose24[i, 0], msg.Pose24[i, 1], msg.Pose24[i, 2], msg.Pose24[i, 3])) * _defaultBindRotation[i];
            }
        }
    }
}