using Sandbox;

/// <summary>
/// Keeps a helper GameObject attached to a named bone on a first-person weapon model.
///
/// Runtime viewmodels have a special ordering constraint: animation bones are settled
/// before the camera modifier chain, while BaseCombatWeapon.PlaceViewModel moves the
/// viewmodel root during that camera chain. Reading a cached world-space bone transform
/// after the root has moved can therefore be one root-transform behind.
///
/// To avoid that, OneMoreRoundWeapon captures each helper as a ROOT-RELATIVE pose before
/// BaseCombatWeapon repositions the viewmodel, then reapplies that same pose after final
/// viewmodel placement. The bone animation remains current, but camera/root movement can
/// no longer create a one-frame feedback loop.
///
/// In editor scenes and as a fallback, this component still follows the final world-space
/// bone transform directly so offsets remain easy to author visually.
/// </summary>
[Title( "OMR Viewmodel Bone Anchor" )]
public sealed class OMRViewModelBoneAnchor : Component, Component.ExecuteInEditor
{
	public readonly struct RootRelativePose
	{
		public readonly Vector3 Position;
		public readonly Rotation Rotation;

		public RootRelativePose( Vector3 position, Rotation rotation )
		{
			Position = position;
			Rotation = rotation;
		}
	}

	[Property, Group( "Reference" )]
	public SkinnedModelRenderer Renderer { get; set; }

	[Property, Group( "Reference" )]
	public string BoneName { get; set; } = "weapon_root";

	[Property, Group( "Offset" )]
	[Description( "Position relative to the selected bone, in model/local units." )]
	public Vector3 PositionOffset { get; set; } = Vector3.Zero;

	[Property, Group( "Offset" )]
	[Description( "Rotation relative to the selected bone." )]
	public Angles RotationOffset { get; set; } = Angles.Zero;

	[Property, Group( "Following" )]
	public bool FollowPosition { get; set; } = true;

	[Property, Group( "Following" )]
	public bool FollowRotation { get; set; } = true;

	[Property, Group( "Diagnostics" )]
	public bool LogMissingBone { get; set; } = false;

	private bool _loggedMissingBone;
	private bool _externallyDrivenThisFrame;

	protected override void OnStart()
	{
		ResolveRenderer();

		// Runtime active viewmodels are driven from OneMoreRoundWeapon.PlaceViewModel.
		// Editor authoring still wants an immediate visual placement.
		if ( Game.IsEditor )
			UpdateAnchorNow();
	}

	protected override void OnUpdate()
	{
		if ( Game.IsEditor )
			UpdateAnchorNow();
	}

	protected override void OnPreRender()
	{
		// The active weapon already supplied a current root-relative transform during
		// its PlaceViewModel stage. Do not overwrite it with a cached world-space bone
		// read later in the same frame.
		if ( _externallyDrivenThisFrame )
		{
			_externallyDrivenThisFrame = false;
			return;
		}

		UpdateAnchorNow();
	}

	/// <summary>
	/// Capture this helper's current animated pose relative to the supplied viewmodel root.
	/// Call this BEFORE BaseCombatWeapon.PlaceViewModel changes that root for the frame.
	/// The old world root cancels out, leaving only the current animated bone pose + authored
	/// helper offset.
	/// </summary>
	public bool TryCaptureRootRelative(
		GameObject viewModelRoot,
		out RootRelativePose pose
	)
	{
		pose = default;
		ResolveRenderer();

		if (
			viewModelRoot is null ||
			!viewModelRoot.IsValid ||
			Renderer is null ||
			!Renderer.IsValid ||
			string.IsNullOrWhiteSpace( BoneName )
		)
		{
			return false;
		}

		if ( !Renderer.TryGetBoneTransform( BoneName, out Transform boneTransform ) )
		{
			LogMissingBoneOnce();
			return false;
		}

		_loggedMissingBone = false;

		Vector3 anchorWorldPosition = boneTransform.PointToWorld( PositionOffset );
		Rotation anchorWorldRotation = boneTransform.RotationToWorld( RotationOffset.ToRotation() );

		Transform rootWorld = viewModelRoot.WorldTransform;
		Vector3 rootLocalPosition = rootWorld.PointToLocal( anchorWorldPosition );
		Rotation rootLocalRotation = rootWorld.RotationToLocal( anchorWorldRotation );

		pose = new RootRelativePose(
			rootLocalPosition,
			rootLocalRotation
		);

		return true;
	}

	/// <summary>
	/// Reapply a captured animated pose after the viewmodel root reaches its final camera
	/// transform for the frame. Marks the helper as externally driven so its fallback
	/// OnPreRender path cannot overwrite the result with stale world-space bone data.
	/// </summary>
	public void ApplyRootRelative(
		GameObject viewModelRoot,
		in RootRelativePose pose
	)
	{
		if ( viewModelRoot is null || !viewModelRoot.IsValid )
			return;

		Transform rootWorld = viewModelRoot.WorldTransform;

		if ( FollowPosition )
			GameObject.WorldPosition = rootWorld.PointToWorld( pose.Position );

		if ( FollowRotation )
			GameObject.WorldRotation = rootWorld.RotationToWorld( pose.Rotation );

		_externallyDrivenThisFrame = true;
	}

	/// <summary>
	/// Direct world-space fallback used for editor authoring and non-weapon contexts.
	/// Runtime active OMR viewmodels normally use the root-relative capture/apply path.
	/// </summary>
	public bool UpdateAnchorNow()
	{
		ResolveRenderer();

		if (
			Renderer is null ||
			!Renderer.IsValid ||
			string.IsNullOrWhiteSpace( BoneName )
		)
		{
			return false;
		}

		if ( !Renderer.TryGetBoneTransform( BoneName, out Transform boneTransform ) )
		{
			LogMissingBoneOnce();
			return false;
		}

		_loggedMissingBone = false;

		Vector3 worldPosition = boneTransform.PointToWorld( PositionOffset );
		Rotation worldRotation = boneTransform.RotationToWorld( RotationOffset.ToRotation() );

		if ( FollowPosition )
			GameObject.WorldPosition = worldPosition;

		if ( FollowRotation )
			GameObject.WorldRotation = worldRotation;

		return true;
	}

	private void ResolveRenderer()
	{
		if ( Renderer is not null && Renderer.IsValid )
			return;

		Renderer = GameObject.Components.Get<SkinnedModelRenderer>(
			FindMode.InAncestors |
			FindMode.InSelf |
			FindMode.InDescendants
		);
	}

	private void LogMissingBoneOnce()
	{
		if ( !LogMissingBone || _loggedMissingBone )
			return;

		_loggedMissingBone = true;
		Log.Warning(
			$"[OMR VM ANCHOR] Bone '{BoneName}' was not found on '{Renderer?.GameObject?.Name ?? "UNKNOWN"}'."
		);
	}
}