#include "OpenXrHandsAdapter.h"

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

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

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

#include "Engine/Engine.h"
#include "IXRTrackingSystem.h"
#include "IXRCamera.h"
#include "Runtime/Launch/Resources/Version.h"
#include "Engine/World.h"
#include "GameFramework/Actor.h"
#include "Camera/CameraComponent.h"
#include "Engine/LocalPlayer.h"

// this is the best docs I found on the subject:
//      https://developer.vive.com/resources/openxr/openxr-pcvr/tutorials/unreal-engine/integrate-hand-tracking-data-your-hand-model/



// Sets default values
AOpenXrHandsAdapter::AOpenXrHandsAdapter()
{

    //UCameraAnchorComponent* 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;
}

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

    _boneTypes = AOpenXrHandsAdapter::GetBoneTypes();

    //HandsAdapterFactory::Register(AR51::OpenXrTethered, this);
    //HandsAdapterFactory::Register(AR51::HtcVive, this);
    HandsAdapterFactory::RegisterAllPlatforms(this);

    TArray<UPoseableMeshComponent*> hands;
    this->GetComponents(hands, false);
    if (hands.Num() > 1) {
        for (UPoseableMeshComponent* HandComponent : hands)
        {
            if (HandComponent)
            {
                FString ComponentName = HandComponent->GetName();
                if (ComponentName.StartsWith("L"))
                {
                    _leftHand = HandComponent;
                }
                else if (ComponentName.StartsWith("R"))
                {
                    _rightHand = HandComponent;
                }
            }
        }
    }
}


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

#if UE_VERSION_OLDER_THAN(5, 7, 0)
    UHeadMountedDisplayFunctionLibrary::GetMotionControllerData(this, EControllerHand::Left, _LHandMotionControllerData);
    UHeadMountedDisplayFunctionLibrary::GetMotionControllerData(this, EControllerHand::Right, _RHandMotionControllerData);
#else
    // UE 5.7+: GetMotionControllerData was removed; hand tracking moved to IXRTrackingSystem.
    // UnrealWorldSpace matches the old function's world-space output.
    if (GEngine && GEngine->XRSystem.IsValid())
    {
        GEngine->XRSystem->GetHandTrackingState(this, EXRSpaceType::UnrealWorldSpace, EControllerHand::Left, _LHandMotionControllerData);
        GEngine->XRSystem->GetHandTrackingState(this, EXRSpaceType::UnrealWorldSpace, EControllerHand::Right, _RHandMotionControllerData);
    }
#endif

    if (ShowHands && _leftHand)
    {
        bool isLHandDisplayed = RenderHands(false);
        _leftHand->SetVisibility(isLHandDisplayed);
    }

    if (ShowHands && _rightHand)
    {
        bool isRHandDisplayed = RenderHands(true);
        _rightHand->SetVisibility(isRHandDisplayed);
    }
}
const TArray<EHandKeypoint>& AOpenXrHandsAdapter::GetBoneTypes() {

    static TArray<EHandKeypoint> boneNames = {

    EHandKeypoint::Wrist,

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

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

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

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

    EHandKeypoint::LittleMetacarpal,
    EHandKeypoint::LittleProximal,
    EHandKeypoint::LittleIntermediate,
    EHandKeypoint::LittleDistal,
    EHandKeypoint::LittleTip

    };
    return boneNames;
}

FString AOpenXrHandsAdapter::GetSourceDevice() const { return "OpenXR"; }

void AOpenXrHandsAdapter::ApplyMotionDataToMesh(const FAR51HandTrackingData& MotionData, bool isRightHand)
{
    UPoseableMeshComponent* MeshComponent;
    if (isRightHand)
        MeshComponent = _rightHand;
    else
        MeshComponent = _leftHand;

    if (MeshComponent)
    {
        // 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
            auto NormalizedRotator = MotionData.HandKeyRotations[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(),
                MotionData.AR51_HAND_KEY_POSITIONS[BoneIndex],
                FVector::OneVector // Scale
            );

            MeshComponent->SetBoneTransformByName(BoneName, NewBoneTransform, EBoneSpaces::WorldSpace);
        }

        // Apply changes to the poseable mesh
        MeshComponent->RefreshBoneTransforms();
    }

}

bool AOpenXrHandsAdapter::RenderHands(bool isRightHand)
{
    bool isHandDisplayed = false;
    auto MotionControllerData = isRightHand ? _RHandMotionControllerData : _LHandMotionControllerData;

    //if (MotionControllerData.bValid && MotionControllerData.TrackingStatus == ETrackingStatus::Tracked && MotionControllerData.HandKeyPositions.Num() > 0)
    if (MotionControllerData.bValid && MotionControllerData.AR51_HAND_KEY_POSITIONS.Num() > 0)
    {
        ApplyMotionDataToMesh(MotionControllerData, isRightHand);
#if UE_VERSION_OLDER_THAN(5, 7, 0)
        UXRVisualizationFunctionLibrary::RenderMotionController(MotionControllerData, isRightHand);
#else
        UXRVisualizationFunctionLibrary::RenderHandTracking(MotionControllerData);
#endif
        isHandDisplayed = true;
    }

    return isHandDisplayed;
}

bool AOpenXrHandsAdapter::GetNativeHmdTransform(FTransform& OutTransform)
{
    const auto& xrSystem = GEngine->XRSystem;
    if (!xrSystem.IsValid())
    {
        return false;
    }

    FQuat Rot;
    auto xrCamera = xrSystem->GetXRCamera();
    //UE_LOG(VAUtilsLog, Log, TEXT("T2W: %s"), *xrSystem->GetTrackingToWorldTransform().ToString());

    //Oculus asserts if you call IsHeadtrackingAllowed not on the render thread, so if we are on the render thread we just assume head tracking is allowed
    const bool headTrackingAllowed = !IsInGameThread() || xrSystem->IsHeadTrackingAllowed();
    FVector Position;

    if (xrCamera.IsValid() &&
        headTrackingAllowed &&
#if ENGINE_MAJOR_VERSION >= 5 && ENGINE_MINOR_VERSION >= 5
        xrCamera->UpdatePlayerCamera(Rot, Position, 0))
#else
        xrCamera->UpdatePlayerCamera(Rot, Position))
#endif
    {
        OutTransform = FTransform(Rot, Position);
        return true;
    }
    return false;
}

UCameraComponent* AOpenXrHandsAdapter::GetMainCameraComponent() const
{
    if (const ULocalPlayer* XRPlayer = GEngine->GetFirstGamePlayer(GetWorld()))
    {
        const APlayerController* PlayerController = XRPlayer->GetPlayerController(GetWorld());
        if (PlayerController)
        {
            const AActor* PlayerViewTarget = PlayerController->GetViewTarget();
            TArray<UCameraComponent*> CameraComponents;
            PlayerViewTarget->GetComponents(CameraComponents);
            UCameraComponent* PlayerCamera = nullptr;
            for (UCameraComponent* Cam : CameraComponents)
            {
                // emulates AActor::CalcCamera(); the PlayerCameraManager uses ViewTarget->CalcCamera() for 
                // the HMD view's basis (regardless of whether bLockToHmd is set), CalcCamera() just finds 
                // the first active cam component and chooses that
                if (Cam->IsActive())
                {
                    return Cam;
                }
            }
        }
    }
    return nullptr;
}

TArray<HandJointInfo> AOpenXrHandsAdapter::GetJointInfos(EHandType handType)
{
    // docs can be found here: https://hub.vive.com/storage/docs/en-us/UnrealPlugin/UnrealHandBlueprints.html
    //   https://developer.vive.com/resources/openxr/openxr-pcvr/tutorials/unreal-engine/integrate-hand-tracking-data-your-hand-model/
    TArray<HandJointInfo> jointInfo;
    bool isRightHand = handType == EHandType::RightHand;
    //auto hand = isRightHand ? EControllerHand::Right : EControllerHand::Left;

    auto MotionControllerData = isRightHand ? _RHandMotionControllerData : _LHandMotionControllerData;

    if (MotionControllerData.bValid && MotionControllerData.AR51_HAND_KEY_POSITIONS.Num() > 0)
    {
        //TTransform root(this);
        //auto w2l = root.GetWorldToLocalMatrix();

        auto fingerPositions = MotionControllerData.AR51_HAND_KEY_POSITIONS;
        auto fingerRotations = MotionControllerData.HandKeyRotations;

        auto distanceBase = fingerPositions[(int)EHandKeypoint::IndexProximal] - fingerPositions[(int)EHandKeypoint::LittleProximal];
        auto palmSize = FMath::Sqrt(distanceBase.X * distanceBase.X + distanceBase.Y * distanceBase.Y + distanceBase.Z * distanceBase.Z);

        if (palmSize > ValidHandMinimumPalmSize)
        {
            FTransform NativeCamera;
            auto cam = GetMainCameraComponent();

            if (!GEngine || !GEngine->XRSystem || !GetNativeHmdTransform(NativeCamera) || !IsValid(cam))
            {
                return jointInfo;
            }

            //const auto TrackingToWorld = GEngine->XRSystem->GetTrackingToWorldTransform(); // CameraParentToWorld
            //const auto WorldToTracking = TrackingToWorld.Inverse(); // WorldToCameraParent
            //const auto ActualCameraToWorld = cam->GetComponentTransform(); // CamereToWorld

            ////const auto CorrectTrackingToWorld = NativeCamera.Inverse()* ActualCameraToWorld;
            //const auto CorrectKpToWorld = WorldToTracking * (NativeCamera.Inverse() * ActualCameraToWorld); // WorldToCameraParent -> (HMD^-1 -> CameraToWorld)

            FTransform WorldToTracking;
            if (HandReferenceSpace == EHandReferenceSpace::HRS_HmdSpace)
            {
                const auto WorldToCamera = cam->GetComponentTransform().Inverse();
                const auto& CameraToTracking = NativeCamera;   // HMD,   NativeCamera is the trasnform of the HMD in the tracking space as reported by the sensors
                WorldToTracking = WorldToCamera * CameraToTracking;
            }
            else if (HandReferenceSpace == EHandReferenceSpace::HRS_CameraRigSpace)
            {
                const auto CameraParentToWorld = GEngine->XRSystem->GetTrackingToWorldTransform(); //usually this is similar to TrackingToWorld
                WorldToTracking = CameraParentToWorld.Inverse(); // WorldToCameraParent
            }

            for (auto i = 0; i < _boneTypes.Num(); i++) {
                auto jointType = (EHandJointType)i;
                EHandKeypoint openXrBoneType = _boneTypes[i];
                int openXrJointIndex = (int)openXrBoneType;

                //FVector p = FVector::ZeroVector;;
                //auto isValid = GetHandJointPosition(hand, openXrJointIndex, p);

                auto p = fingerPositions[openXrJointIndex];
                auto r = fingerRotations[openXrJointIndex];
                auto confidence = 1.f;

                auto lp = WorldToTracking.TransformPosition(p);
                auto lr = WorldToTracking.GetRotation() * r;

                //p = CorrectKpToWorld.TransformPosition(p);
                //r = CorrectKpToWorld.GetRotation() * r;

                jointInfo.Add({ handType, jointType, confidence, lp, lr, p, r });
            }
        }
        else {
            LOG("Debug: isRightHand=%d, palm size is too small (%f) does not tramitting hands", isRightHand, (float)palmSize);
        }
    }

    return jointInfo;
}

TArray<HandJointInfo> AOpenXrHandsAdapter::GetRightJointInfos() { return GetJointInfos(EHandType::RightHand); }

TArray<HandJointInfo> AOpenXrHandsAdapter::GetLeftJointInfos() { return GetJointInfos(EHandType::LeftHand); }

