// Copyright 2023 PICO Inc. All Rights Reserved.

#include "PICO_BodyTracking.h"
#include "PICO_MovementModule.h"
#include "OpenXRCore.h"
#include "PICOOpenXRRuntimeSettings.h"

#if PLATFORM_ANDROID
#include "Android/AndroidApplication.h"
#endif

FBodyTrackingPICO::FBodyTrackingPICO()
{
	Locations.SetNum(XR_BODY_JOINT_COUNT_BD);
}

void FBodyTrackingPICO::Register()
{
	RegisterOpenXRExtensionModularFeature();
}

void FBodyTrackingPICO::Unregister()
{
	UnregisterOpenXRExtensionModularFeature();
}

bool FBodyTrackingPICO::GetRequiredExtensions(TArray<const ANSICHAR*>& OutExtensions)
{
	if (UPICOOpenXRRuntimeSettings::GetBoolConfigByKey("bEnableBodyTracking"))
	{
		OutExtensions.Add("XR_BD_body_tracking");
		return true;
	}
	return false;
}

const void* FBodyTrackingPICO::OnGetSystem(XrInstance InInstance, const void* InNext)
{
	XR_ENSURE(xrGetInstanceProcAddr(InInstance, "xrCreateBodyTrackerBD", (PFN_xrVoidFunction*)&xrCreateBodyTrackerBD));
	XR_ENSURE(xrGetInstanceProcAddr(InInstance, "xrDestroyBodyTrackerBD", (PFN_xrVoidFunction*)&xrDestroyBodyTrackerBD));
	XR_ENSURE(xrGetInstanceProcAddr(InInstance, "xrLocateBodyJointsBD", (PFN_xrVoidFunction*)&xrLocateBodyJointsBD));
	Instance = InInstance;
	return InNext;
}

void FBodyTrackingPICO::PostGetSystem(XrInstance InInstance, XrSystemId InSystem)
{
	XrSystemBodyTrackingPropertiesBD BodyTrackingSystemProperties = { (XrStructureType)XR_TYPE_SYSTEM_BODY_TRACKING_PROPERTIES_BD };
	XrSystemProperties systemProperties{ XR_TYPE_SYSTEM_PROPERTIES,&BodyTrackingSystemProperties };
	XR_ENSURE(xrGetSystemProperties(InInstance, InSystem, &systemProperties));
	bCurrentDeviceSupportBodyTracking = BodyTrackingSystemProperties.supportsBodyTracking == XR_TRUE;
}

void FBodyTrackingPICO::PostCreateSession(XrSession InSession)
{
	Session = InSession;
}

void FBodyTrackingPICO::UpdateDeviceLocations(XrSession InSession, XrTime DisplayTime, XrSpace TrackingSpace)
{
	PredictedTime = DisplayTime;
	BaseSpace = TrackingSpace;
}

bool FBodyTrackingPICO::TryGetBodyState(FBodyStatePICO& outBodyState, float WorldToMeters)
{
	if (bBodyTrackerIsRunning)
	{
		XrBodyJointsLocateInfoBD GetInfo = { (XrStructureType)XR_TYPE_BODY_JOINTS_LOCATE_INFO_BD };
		GetInfo.baseSpace = BaseSpace;
		GetInfo.time = PredictedTime;

		outBodyState.IsActive = false;
		outBodyState.BaseJointsData.SetNum(Locations.Num());
		XrBodyJointLocationsBD BodyData = { (XrStructureType)XR_TYPE_BODY_JOINT_LOCATIONS_BD };
		BodyData.jointLocationCount = Locations.Num();
		BodyData.jointLocations = Locations.GetData();

		if (XR_SUCCEEDED(xrLocateBodyJointsBD(BodyTracker, &GetInfo, &BodyData)))
		{
			outBodyState.IsActive = BodyData.allJointPosesTracked == XR_TRUE;

			if (outBodyState.IsActive)
			{
				for (int i = 0; i < XR_BODY_JOINT_COUNT_BD; ++i)
				{
					outBodyState.BaseJointsData[i].Joint = (EBodyJointPICO)i;
					outBodyState.BaseJointsData[i].bIsValid = Locations[i].locationFlags & (XR_SPACE_LOCATION_ORIENTATION_VALID_BIT | XR_SPACE_LOCATION_POSITION_VALID_BIT);
					outBodyState.BaseJointsData[i].Orientation = FRotator(ToFQuat(Locations[i].pose.orientation));
					outBodyState.BaseJointsData[i].Position = ToFVector(Locations[i].pose.position, WorldToMeters);
				}
			}
			return true;
		}
	}

	outBodyState.IsActive = false;
	return false;
}

bool FBodyTrackingPICO::IsBodyTrackingEnabled()
{
	return bBodyTrackerIsRunning;
}

bool FBodyTrackingPICO::IsBodyTrackingSupported()
{
	return bCurrentDeviceSupportBodyTracking;
}

bool FBodyTrackingPICO::StartBodyTracking(EBodyTrackingModePICO Mode)
{
	if (bCurrentDeviceSupportBodyTracking && Session != XR_NULL_HANDLE && BodyTracker == XR_NULL_HANDLE)
	{
		XrBodyTrackerCreateInfoBD CreateInfo = { (XrStructureType)XR_TYPE_BODY_TRACKER_CREATE_INFO_BD };
		CreateInfo.jointSet = (XrBodyJointSetBD)Mode;
		if (XR_FAILED(xrCreateBodyTrackerBD(Session, &CreateInfo, &BodyTracker)))
		{
			return false;
		}
		UE_LOG(PICOOpenXRMovement, Log, TEXT("Created BodyTracker."));
		bBodyTrackerIsRunning = true;
		return true;
	}
	return false;
}

bool FBodyTrackingPICO::StopBodyTracking()
{
	if (BodyTracker != XR_NULL_HANDLE)
	{
		UE_LOG(PICOOpenXRMovement, Log, TEXT("Destroyed BodyTracker."));
		bBodyTrackerIsRunning = false;
		XrResult Result = xrDestroyBodyTrackerBD(BodyTracker);
		if (XR_SUCCEEDED(Result))
		{
			BodyTracker = XR_NULL_HANDLE;
			return true;
		}
		return false;
	}
	return false;
}