using Sandbox;
using System;
using System.Collections.Generic;

/// <summary>
/// Final owner-local viewmodel placement extensions.
///
/// Frame ordering is deliberate:
/// 1) capture every bone helper ROOT-RELATIVE while current animation bones are settled;
/// 2) let BaseCombatWeapon pin the viewmodel root to the composed CameraView;
/// 3) center the authored optic using the captured root-relative lens pose;
/// 4) reapply every helper from the final root before PreRender.
///
/// This removes both sources of one-frame lag from the previous implementation:
/// the alignment no longer chases a cached world-space bone after moving the root, and
/// WorldPanel / PiP helpers no longer wait until an unordered OnPreRender callback to
/// learn where this frame's scope actually is.
/// </summary>
public partial class OneMoreRoundWeapon
{
	private readonly List<OMRViewModelBoneAnchor> _cachedViewModelBoneAnchors = new();
	private readonly List<CapturedViewModelBoneAnchor> _capturedViewModelBoneAnchors = new();
	private GameObject _boneAnchorCacheViewModel;

	private readonly struct CapturedViewModelBoneAnchor
	{
		public readonly OMRViewModelBoneAnchor Anchor;
		public readonly OMRViewModelBoneAnchor.RootRelativePose Pose;
		public readonly bool IsValid;

		public CapturedViewModelBoneAnchor(
			OMRViewModelBoneAnchor anchor,
			in OMRViewModelBoneAnchor.RootRelativePose pose,
			bool isValid
		)
		{
			Anchor = anchor;
			Pose = pose;
			IsValid = isValid;
		}
	}

	protected override void PlaceViewModel(
		CameraComponent camera,
		in CameraView view
	)
	{
		GameObject viewModel = ViewModel;

		if ( viewModel is null || !viewModel.IsValid )
		{
			base.PlaceViewModel( camera, view );
			ClearViewModelBoneAnchorCache();
			return;
		}

		EnsureViewModelBoneAnchorCache( viewModel );
		CaptureViewModelBoneAnchors( viewModel );

		// Canonical engine placement first.
		base.PlaceViewModel( camera, view );

		OMRViewModelAimAlignment alignment =
			viewModel.Components.Get<OMRViewModelAimAlignment>();

		if ( alignment is not null && alignment.IsValid )
		{
			OMROpticProfile profile = ActiveOpticProfile;

			if (
				profile is not null &&
				profile.IsValid &&
				profile.EnableViewModelAimAlignment &&
				TryGetCapturedAnchor(
					profile.ViewModelAimAnchorName,
					out OMRViewModelBoneAnchor.RootRelativePose alignmentPose
				)
			)
			{
				alignment.ApplyAlignment(
					this,
					camera,
					view,
					alignmentPose
				);
			}
		}

		// The root is now final for this frame. Publish every helper immediately so
		// the PiP camera and WorldPanel see the same transform before PreRender.
		ApplyCapturedViewModelBoneAnchors( viewModel );
	}

	private void EnsureViewModelBoneAnchorCache( GameObject viewModel )
	{
		if ( _boneAnchorCacheViewModel == viewModel )
			return;

		_boneAnchorCacheViewModel = viewModel;
		_cachedViewModelBoneAnchors.Clear();
		_capturedViewModelBoneAnchors.Clear();

		foreach ( GameObject obj in viewModel.GetAllObjects( true ) )
		{
			if ( obj is null || !obj.IsValid )
				continue;

			OMRViewModelBoneAnchor anchor =
				obj.GetComponent<OMRViewModelBoneAnchor>( true );

			if ( anchor is null || !anchor.IsValid )
				continue;

			_cachedViewModelBoneAnchors.Add( anchor );
		}
	}

	private void CaptureViewModelBoneAnchors( GameObject viewModel )
	{
		_capturedViewModelBoneAnchors.Clear();

		for ( int i = 0; i < _cachedViewModelBoneAnchors.Count; i++ )
		{
			OMRViewModelBoneAnchor anchor = _cachedViewModelBoneAnchors[i];

			if ( anchor is null || !anchor.IsValid )
				continue;

			bool captured = anchor.TryCaptureRootRelative(
				viewModel,
				out OMRViewModelBoneAnchor.RootRelativePose pose
			);

			_capturedViewModelBoneAnchors.Add(
				new CapturedViewModelBoneAnchor(
					anchor,
					pose,
					captured
				)
			);
		}
	}

	private bool TryGetCapturedAnchor(
		string objectName,
		out OMRViewModelBoneAnchor.RootRelativePose pose
	)
	{
		pose = default;

		if ( string.IsNullOrWhiteSpace( objectName ) )
			return false;

		for ( int i = 0; i < _capturedViewModelBoneAnchors.Count; i++ )
		{
			CapturedViewModelBoneAnchor captured =
				_capturedViewModelBoneAnchors[i];

			if (
				!captured.IsValid ||
				captured.Anchor is null ||
				!captured.Anchor.IsValid ||
				!string.Equals(
					captured.Anchor.GameObject.Name,
					objectName,
					StringComparison.OrdinalIgnoreCase
				)
			)
			{
				continue;
			}

			pose = captured.Pose;
			return true;
		}

		return false;
	}

	private void ApplyCapturedViewModelBoneAnchors( GameObject viewModel )
	{
		for ( int i = 0; i < _capturedViewModelBoneAnchors.Count; i++ )
		{
			CapturedViewModelBoneAnchor captured =
				_capturedViewModelBoneAnchors[i];

			if (
				!captured.IsValid ||
				captured.Anchor is null ||
				!captured.Anchor.IsValid
			)
			{
				continue;
			}

			captured.Anchor.ApplyRootRelative(
				viewModel,
				captured.Pose
			);
		}
	}

	private void ClearViewModelBoneAnchorCache()
	{
		_boneAnchorCacheViewModel = null;
		_cachedViewModelBoneAnchors.Clear();
		_capturedViewModelBoneAnchors.Clear();
	}
}