#include "PicoHandsAdapter.h"

#include "HandsAdapterFactory.h"
#include "BoneTransform.h"

#include "TTransform.h"
#include "Logging/LogMacros.h"
#include "HeadMountedDisplay.h"

#include "PXR_HMDRuntimeSettings.h"
#include "PXR_InputFunctionLibrary.h"
//#include <UnityAdapter.h>

//#include <HeadMountedDisplayFunctionLibrary.h>
#include <MotionControllerComponent.h>

// this is the best docs I found on the subject:
//      https://developer.picoxr.com/document/unreal/hand-tracking/

CPicoHand::CPicoHand()
{
    Init(EPICOXRHandType::None, nullptr);
}

CPicoHand::CPicoHand(EPICOXRHandType eHandType, AActor* pOwner)
{
    Init(eHandType, pOwner);
}


CPicoHand::~CPicoHand()
{

}

bool CPicoHand::Init(EPICOXRHandType eHandType, AActor* pOwner)
{
    _handType = eHandType;
    _pOwner = pOwner;
    _bVisible = false;
    _handMesh = nullptr;
    _handSkinnedAsset = nullptr;
    _handMaterial = nullptr;

    _jointPositions.Init(FVector::ZeroVector, static_cast<int32>(EHandJointCount));
    _jointRotations.Init(FQuat::Identity, static_cast<int32>(EHandJointCount));

    for (uint8 BoneIndex = 0; BoneIndex < static_cast<uint8>(EHandJointCount); BoneIndex++)
    {
        EPICOXRHandJoint picoXRJoint = (EPICOXRHandJoint)BoneIndex;
        _boneNameMappings.Add(picoXRJoint,
                              FormatHandJointWithPrefix((eHandType == EPICOXRHandType::HandLeft) ? TEXT("left"): ((eHandType == EPICOXRHandType::HandRight) ? TEXT("right"): TEXT("none")),
                              picoXRJoint));
    }

    return true;
}

FName CPicoHand::FormatHandJointWithPrefix(const FString& Prefix, EPICOXRHandJoint Joint)
{
    const UEnum* EnumPtr = StaticEnum<EPICOXRHandJoint>();
    if (!EnumPtr)
    {
        return FName(*Prefix.ToLower());
    }

    FString EnumName = EnumPtr->GetNameStringByValue(static_cast<int64>(Joint));

    FString Result = Prefix.ToLower();
    for (int32 i = 0; i < EnumName.Len(); ++i)
    {
        const TCHAR CurrentChar = EnumName[i];

        if (FChar::IsUpper(CurrentChar))
        {
            Result += "_";
            Result += FChar::ToLower(CurrentChar);
        }
        else
        {
            Result += CurrentChar;
        }
    }

    return FName(*Result);
}

void CPicoHand::OnBeginPlay(USkinnedAsset* skinnedAsset, UMaterial* material)
{
    if (_handMesh)
    {
        _handMesh->DestroyComponent();
    }

    _handSkinnedAsset = skinnedAsset;
    _handMaterial = material;

    // Create an instance of UPoseableMeshComponent from RFNightHandSkinnedAsset and store it in _rightHand
    if (_handSkinnedAsset)
    {
        _handMesh = NewObject<UPoseableMeshComponent>(_pOwner, _handType == EPICOXRHandType::HandLeft ? TEXT("LeftHandMeshComponent"): TEXT("RightHandMeshComponent"));
        if (_handMesh)
        {
            _handMesh->SetSkeletalMesh(Cast<USkeletalMesh>(_handSkinnedAsset));
            _handMesh->SetupAttachment(_pOwner->GetRootComponent());
            _handMesh->RegisterComponent();

            if (_handMaterial != NULL)
            {
                _handMesh->SetMaterial(0, _handMaterial);
            }
        }
    }

    _boneTypes = CPicoHand::GetBoneTypes();
}

void CPicoHand::UpdateHandData()
{
    GetHandTrackingDataFromPicoSDK();

    FTransform HandPoseRight = UPICOXRInputFunctionLibrary::GetHandRootPose(_handType);
    if (HandPoseRight.IsValid() && !HandPoseRight.Equals(FTransform()))
    {
        if (!HandPoseRight.GetLocation().ContainsNaN())
            _jointPositions[(int)EPICOXRHandJoint::Wrist] = HandPoseRight.GetLocation();
    }
}

bool CPicoHand::UpdateVisibility(bool showHands, bool bApplyLocationToBones, bool hideByConfidence,
                                 EPICOXRHandTrackingConfidence confidenceThreshold)
{
    if (showHands)
    {
        // Update the visible pose meshes if ShowHands is true.
        if (_handMesh)
        {
            bool handDataValid = (_jointPositions.Num() > 0) && (_jointRotations.Num() > 0);
            if (hideByConfidence)
                handDataValid &= _trackingConfidence == EPICOXRHandTrackingConfidence::High;

            if (_handSkinnedAsset != NULL && handDataValid)
            {
                UpdateBonePose(bApplyLocationToBones);
                UpdateHandTransform();
            }

            SetVisibility(handDataValid);
        }
    }
    else 
    {
        SetVisibility(false);
    }

    return _bVisible;
}

void CPicoHand::UpdateBonePose(bool bApplyLocationToBones)
{
    if ((_handSkinnedAsset != NULL) && (_handMesh != NULL))
    {
        for (uint8 BoneIndex = 0; BoneIndex < static_cast<uint8>(EHandJointCount); BoneIndex++)
        {
            EPICOXRHandJoint picoXRJoint = (EPICOXRHandJoint)BoneIndex;
            FName boneName = _boneNameMappings.FindRef(picoXRJoint);

            FQuat BoneRotation = _jointRotations[BoneIndex];
            if (!BoneRotation.IsIdentity() && BoneRotation.IsNormalized())
            {
                _handMesh->SetBoneRotationByName(boneName, BoneRotation.Rotator(), EBoneSpaces::ComponentSpace);
            }

            const FVector BoneLocation = _jointPositions[BoneIndex];
            if (!BoneLocation.IsZero() && !BoneLocation.ContainsNaN() && bApplyLocationToBones)
            {
                _handMesh->SetBoneLocationByName(boneName, BoneLocation, EBoneSpaces::WorldSpace);
            }

        }
    }
}

void CPicoHand::UpdateHandTransform()
{
    //pPosableMeshComponent.SetHiddenInGame(true, true);
    if (_handMesh)
        _handMesh->SetRelativeLocation(_jointPositions[(int)EPICOXRHandJoint::Wrist]);
}

void CPicoHand::SetVisibility(bool bNewVisiblity)
{
    if (_handMesh && (_bVisible != bNewVisiblity))
    {
        _handMesh->SetVisibility(bNewVisiblity);
    }

    _bVisible = bNewVisiblity;
}

void CPicoHand::GetHandTrackingDataFromPicoSDK()
{
    // Update the hand joint positions and rotations
    for (uint8 BoneIndex = 0; BoneIndex < static_cast<uint8>(EHandJointCount); BoneIndex++)
    {
        EPICOXRHandJoint picoXRJoint = (EPICOXRHandJoint)BoneIndex;
        _jointPositions[BoneIndex] = UPICOXRInputFunctionLibrary::GetBoneLocation(_handType, picoXRJoint);
        _jointRotations[BoneIndex] = UPICOXRInputFunctionLibrary::GetBoneRotation(_handType, picoXRJoint);
    }

    _trackingConfidence = UPICOXRInputFunctionLibrary::GetTrackingConfidence(_handType);
}

//void CPicoHand::ApplyMotionDataToMesh()
//{
//    if (_handMesh)
//    {
//        // Loop over all bones (joints)
//        auto NumBones = (int32)EHandKeypoint::LittleTip - (int32)EHandKeypoint::Palm;
//        for (int32 BoneIndex = 0; BoneIndex < NumBones; ++BoneIndex)
//        {
//            //FName BoneName = MeshComponent->GetBoneName(BoneIndex);
//            FString BoneString = FString::Printf(TEXT("WaveBone_%d"), BoneIndex);
//            FName BoneName = FName(*BoneString);
//
//            // Set the bone's transform based on MotionData
//            FRotator NormalizedRotator = _jointRotations.operator[](BoneIndex).GetNormalized().Rotator();
//
//            // Convert the normalized rotator to a quaternion
//            FQuat RotatorAsQuat = NormalizedRotator.Quaternion();
//
//            // UE4 and UE5 uses different mashes. 
//#ifdef UE_5_0_OR_LATER
//             // Apply transformations: Flip roll and pitch, and rotate yaw by 180 degrees
//            FQuat Rotation180Yaw(FVector::UpVector, FMath::DegreesToRadians(180.0f));
//            FQuat FlippedQuat = FQuat(FVector(-1.0f, -1.0f, 1.0f), 0.0f);
//            // Combine the transformations
//            FQuat ModifiedQuat = RotatorAsQuat * Rotation180Yaw * FlippedQuat;
//#else
//            FQuat ModifiedQuat = RotatorAsQuat;
//#endif
//            FTransform NewBoneTransform(
//                ModifiedQuat.GetNormalized(),
//                _jointPositions.operator[](BoneIndex),
//                FVector::OneVector // Scale
//            );
//
//            _handMesh->SetBoneTransformByName(BoneName, NewBoneTransform, EBoneSpaces::WorldSpace);
//        }
//
//        // Apply changes to the poseable mesh
//        _handMesh->RefreshBoneTransforms();
//    }
//
//}

TArray<HandJointInfo> CPicoHand::GetJointInfos(float validHandMinimumPalmSize)
{
    // docs can be found here: https://developer.picoxr.com/document/unreal/hand-tracking/

    TArray<HandJointInfo> jointInfo;

    if (UPICOXRInputFunctionLibrary::IsHandTrackingEnabled())
    {
        // Get the root transform of the hand.
        TTransform root(_pOwner);
        auto l2w = root.GetLocalToWorldMatrix();
        auto w2l = root.GetWorldToLocalMatrix();

        EHandType eHandType = (_handType == EPICOXRHandType::HandRight ? EHandType::RightHand : EHandType::LeftHand);
        auto distanceBase = _jointPositions[(int)EPICOXRHandJoint::IndexProximal] - _jointPositions[(int)EPICOXRHandJoint::LittleProximal];
        auto palmSize = FMath::Sqrt(distanceBase.X * distanceBase.X + distanceBase.Y * distanceBase.Y + distanceBase.Z * distanceBase.Z);

        if (palmSize > validHandMinimumPalmSize)
        {
            for (auto i = 0; i < _boneTypes.Num(); i++) {
                auto ar51JointType = (EHandJointType)i; //ar51 joint type
                EPICOXRHandJoint picoJoint = _boneTypes[i];

                // Retrieve bone location and rotation via Pico XR functions.
                FVector lp = _jointPositions[(int)picoJoint];
                FQuat r = _jointRotations[(int)picoJoint];

                // A simple confidence value.
                float confidence = 1.0f;

                // Transform the global position to local space.
                FVector p = l2w.TransformPosition(lp);

                // Transform rotation to local space.
                FQuat lr = w2l.Rotator().Quaternion() * r;

                // Use the joint index cast to EHandJointType.
                jointInfo.Add({ eHandType, ar51JointType, confidence, lp, lr, p, r });
            }
        }
        // Should these be uncommented?
       // else {
       //         LOG("Debug: isRightHand=%d, palm size is too small (%f) does not tramitting hands", isRightHand, (float)palmSize);
       //      }
    }
    /*else
    {
        LOG("Hand tracking disabled or hand component invalid for handType=%d", isRightHand);
    }*/

    return jointInfo;
}

const TArray<EPICOXRHandJoint>& CPicoHand::GetBoneTypes()
{
    static TArray<EPICOXRHandJoint> boneNames = {
        //EPICOXRHandJoint::Palm,
        EPICOXRHandJoint::Wrist,

        EPICOXRHandJoint::ThumbMetacarpal, // note: we place this twice
        EPICOXRHandJoint::ThumbMetacarpal, // note: same value
        EPICOXRHandJoint::ThumbProximal,
        EPICOXRHandJoint::ThumbDistal,
        EPICOXRHandJoint::ThumbTip,

        //EPICOXRHandJoint::IndexMetacarpal,
        EPICOXRHandJoint::IndexProximal,
        EPICOXRHandJoint::IndexIntermediate,
        EPICOXRHandJoint::IndexDistal,
        EPICOXRHandJoint::IndexTip,

        //EPICOXRHandJoint::MiddleMetacarpal,
        EPICOXRHandJoint::MiddleProximal,
        EPICOXRHandJoint::MiddleIntermediate,
        EPICOXRHandJoint::MiddleDistal,
        EPICOXRHandJoint::MiddleTip,

        //EPICOXRHandJoint::RingMetacarpal,
        EPICOXRHandJoint::RingProximal,
        EPICOXRHandJoint::RingIntermediate,
        EPICOXRHandJoint::RingDistal,
        EPICOXRHandJoint::RingTip,

        EPICOXRHandJoint::LittleMetacarpal,
        EPICOXRHandJoint::LittleProximal,
        EPICOXRHandJoint::LittleIntermediate,
        EPICOXRHandJoint::LittleDistal,
        EPICOXRHandJoint::LittleTip

    };
    return boneNames;
}


// Sets default values
APicoHandsAdapter::APicoHandsAdapter()
{
    _cameraAnchorComponent = CreateDefaultSubobject<UCameraAnchorComponent>(TEXT("CameraAnchorSceneComp"));
    SetRootComponent(_cameraAnchorComponent); // Set this component as the root

    // Set this actor to call Tick() every frame.  You can turn this off to improve performance if you don't need it.
    PrimaryActorTick.bCanEverTick = true;

    _pLeftHand = new CPicoHand(EPICOXRHandType::HandLeft, this);
    _pRightHand = new CPicoHand(EPICOXRHandType::HandRight, this);
}

// Called when the game starts or when spawned
void APicoHandsAdapter::BeginPlay()
{
    Super::BeginPlay();

    _pLeftHand->OnBeginPlay(LeftHandSkinnedAsset, _handsMaterial);
    _pRightHand->OnBeginPlay(RightHandSkinnedAsset, _handsMaterial);

    HandsAdapterFactory::RegisterAllPlatforms(this);
}

// Called every frame
void APicoHandsAdapter::Tick(float DeltaTime)
{
    Super::Tick(DeltaTime);

    if (IsInGameThread())
    {
        if (!GetOwner()->HasLocalNetOwner())
            return;
    }
    else
        return;

    if (!UPICOXRInputFunctionLibrary::IsHandTrackingEnabled())
    {
        _pLeftHand->SetVisibility(false);
        _pRightHand->SetVisibility(false);

		return;
    }
    
    _pLeftHand->UpdateHandData();
    _pLeftHand->UpdateVisibility(ShowHands, _bHideByConfidence, bApplyLocationToBones, EPICOXRHandTrackingConfidence::High);

    _pRightHand->UpdateHandData();
    _pRightHand->UpdateVisibility(ShowHands, _bHideByConfidence, bApplyLocationToBones, EPICOXRHandTrackingConfidence::High);

   
}

FString APicoHandsAdapter::GetSourceDevice() const { return "Pico"; }

// docs can be found here: https://developer.picoxr.com/document/unreal/hand-tracking/

TArray<HandJointInfo> APicoHandsAdapter::GetRightJointInfos() { return _pRightHand->GetJointInfos(ValidHandMinimumPalmSize); }

TArray<HandJointInfo> APicoHandsAdapter::GetLeftJointInfos() { return _pLeftHand->GetJointInfos(ValidHandMinimumPalmSize); }

