Vertex Alignment Tool

using Editor;
using Sandbox;
using System.Collections.Generic;
using System.Linq;

namespace Editor
{
	[EditorTool( "tools.vertex-alignment" )]
	[Title( "Vertex Alignment Tool" )]
	[Icon( "control_point" )]
	[Alias( ["vertex-align", "vertex-snap"] )]
	[Group( "2" )]
	public class VertexAlignmentTool : EditorTool
	{
		private enum AlignmentState
		{
			SelectingSourceVertex,
			SelectingTargetVertex
		}

		private AlignmentState currentState = AlignmentState.SelectingSourceVertex;
		private Vector3? sourceVertex = null;
		private GameObject sourceObject = null;
		private readonly List<VertexInfo> sceneVertices = new List<VertexInfo>();
		private readonly int vertexDetectionRadius = 15;

		private struct VertexInfo
		{
			public Vector3 WorldPosition;
			public Vector3 LocalPosition;
			public GameObject GameObject;
			public ModelRenderer Renderer;
		}

		public override void OnUpdate()
		{
			CollectSceneVertices();
			DrawVertices();

			if ( Gizmo.WasLeftMousePressed )
			{
				HandleVertexSelection();
			}

			if ( Application.IsKeyDown( KeyCode.Escape ) )
			{
				ResetState();
			}
		}

		private void ResetState()
		{
			currentState = AlignmentState.SelectingSourceVertex;
			sourceVertex = null;
			sourceObject = null;
		}

		private void CollectSceneVertices()
		{
			sceneVertices.Clear();

			var session = SceneEditorSession.Active;
			if ( session?.Scene == null ) return;

			foreach ( var meshRenderer in session.Scene.GetAllComponents<ModelRenderer>() )
			{
				if ( !meshRenderer.IsValid || meshRenderer.Model == null )
					continue;

				var model = meshRenderer.Model;
				var gameObject = meshRenderer.GameObject;
				var transform = meshRenderer.WorldTransform;

				try
				{
					var vertices = model.GetVertices();
					if ( vertices == null ) continue;

					foreach ( var vertex in vertices )
					{
						var worldPos = transform.PointToWorld( vertex.Position );
						sceneVertices.Add( new VertexInfo
						{
							WorldPosition = worldPos,
							LocalPosition = vertex.Position,
							GameObject = gameObject,
							Renderer = meshRenderer
						} );
					}
				}
				catch
				{
					continue;
				}
			}
		}

		private void DrawVertices()
		{
			if ( sceneVertices.Count == 0 ) return;

			var mousePosition = SceneViewportWidget.MousePosition;
			var closestVertex = GetClosestVertex( mousePosition );

			using ( Gizmo.Scope( "vertex_visualization" ) )
			{
				Gizmo.Transform = Transform.Zero;

				if ( closestVertex.HasValue )
				{
					var centerVertex = closestVertex.Value;
					var centerPos = centerVertex.WorldPosition;
					var cameraPos = Gizmo.Camera.Position;
					var distance = Vector3.DistanceBetween( cameraPos, centerPos );
					var baseSize = MathX.Remap( distance, 10f, 100f, 4f, 1f ).Clamp( 0.5f, 6f );

					Color centerColor;
					if ( currentState == AlignmentState.SelectingSourceVertex )
					{
						centerColor = Color.Yellow;
					}
					else if ( sourceVertex.HasValue && Vector3.DistanceBetween( centerPos, sourceVertex.Value ) < 0.1f )
					{
						centerColor = Color.Green;
					}
					else
					{
						centerColor = Color.Cyan;
					}

					Gizmo.Draw.Color = centerColor.WithAlpha( 1.0f );
					Gizmo.Draw.SolidSphere( centerPos, baseSize * 1.5f );
					Gizmo.Draw.Color = centerColor.WithAlpha( 0.6f );
					Gizmo.Draw.SolidSphere( centerPos, baseSize * 0.8f );

					var nearbyVertices = sceneVertices
						.Where( v => Vector3.DistanceBetween( v.WorldPosition, centerPos ) > baseSize * 3f )
						.Where( v => Vector3.DistanceBetween( v.WorldPosition, centerPos ) < baseSize * 20f )
						.OrderBy( v => Vector3.DistanceBetween( v.WorldPosition, centerPos ) )
						.Take( 8 )
						.ToList();

					foreach ( var vertexInfo in nearbyVertices )
					{
						var vertex = vertexInfo.WorldPosition;
						var vertexDistance = Vector3.DistanceBetween( cameraPos, vertex );
						var size = MathX.Remap( vertexDistance, 10f, 100f, 3f, 0.8f ).Clamp( 0.3f, 4f );

						Color color = Color.Blue;
						if ( currentState == AlignmentState.SelectingTargetVertex &&
							 sourceVertex.HasValue &&
							 Vector3.DistanceBetween( vertex, sourceVertex.Value ) < 0.1f )
						{
							color = Color.Green;
						}

						Gizmo.Draw.Color = color.WithAlpha( 0.6f );
						Gizmo.Draw.SolidSphere( vertex, size );
						Gizmo.Draw.Color = color.WithAlpha( 0.8f );
						Gizmo.Draw.SolidSphere( vertex, size * 0.4f );
					}
				}

				if ( sourceVertex.HasValue && closestVertex.HasValue )
				{
					Gizmo.Draw.Color = Color.White.WithAlpha( 0.5f );
					Gizmo.Draw.LineThickness = 2f;
					Gizmo.Draw.Line( sourceVertex.Value, closestVertex.Value.WorldPosition );
				}
			}
		}

		private List<VertexInfo> GetNearbyVertices( Vector2 mousePosition, int radius )
		{
			var nearbyVertices = new List<VertexInfo>();
			var rays = new List<Ray> { Gizmo.CurrentRay };

			for ( int i = 1; i < radius; i++ )
			{
				rays.Add( Gizmo.Camera.GetRay( mousePosition + new Vector2( 0f, i ) ) );
				rays.Add( Gizmo.Camera.GetRay( mousePosition + new Vector2( i, 0f ) ) );
				rays.Add( Gizmo.Camera.GetRay( mousePosition + new Vector2( 0f, -i ) ) );
				rays.Add( Gizmo.Camera.GetRay( mousePosition + new Vector2( -i, 0f ) ) );
			}

			foreach ( var vertexInfo in sceneVertices )
			{
				foreach ( var ray in rays )
				{
					var distance = DistanceFromPointToRay( ray, vertexInfo.WorldPosition );
					if ( distance < 5f )
					{
						nearbyVertices.Add( vertexInfo );
						break;
					}
				}
			}

			return nearbyVertices;
		}

		private void HandleVertexSelection()
		{
			var closestVertex = GetClosestVertex( SceneViewportWidget.MousePosition );

			if ( !closestVertex.HasValue ) return;

			if ( currentState == AlignmentState.SelectingSourceVertex )
			{
				sourceVertex = closestVertex.Value.WorldPosition;
				sourceObject = closestVertex.Value.GameObject;
				EditorScene.Selection.Set( sourceObject );
				currentState = AlignmentState.SelectingTargetVertex;
			}
			else
			{
				var targetVertex = closestVertex.Value.WorldPosition;
				PerformAlignment( targetVertex );
			}
		}

		private VertexInfo? GetClosestVertex( Vector2 mousePosition )
		{
			var nearbyVertices = GetNearbyVertices( mousePosition, vertexDetectionRadius );

			VertexInfo? closest = null;
			float closestDistance = float.MaxValue;

			foreach ( var vertexInfo in nearbyVertices )
			{
				var distance = DistanceFromPointToRay( Gizmo.CurrentRay, vertexInfo.WorldPosition );

				if ( distance < closestDistance )
				{
					closest = vertexInfo;
					closestDistance = distance;
				}
			}

			return closest;
		}

		private float DistanceFromPointToRay( Ray ray, Vector3 point )
		{
			var toPoint = point - ray.Position;
			var projection = Vector3.Dot( toPoint, ray.Forward );
			var closestPointOnRay = ray.Position + ray.Forward * projection;
			return Vector3.DistanceBetween( point, closestPointOnRay );
		}

		private void PerformAlignment( Vector3 targetVertex )
		{
			if ( !sourceVertex.HasValue || sourceObject == null ) return;

			var session = SceneEditorSession.Active;
			if ( session?.Scene == null ) return;

			var offset = targetVertex - sourceVertex.Value;

			var scope = session.UndoScope( "Align Vertices" ).WithGameObjectChanges( new[] { sourceObject }, GameObjectUndoFlags.Properties );
			using ( scope.Push() )
			{
				sourceObject.WorldPosition += offset;
			}

			ResetState();
		}

		[Shortcut( "tools.vertex-alignment", "V", ShortcutType.Widget )]
		public static void ActivateTool()
		{
			EditorToolManager.SetTool( "VertexAlignmentTool" );
		}
	}
}