Editor/HierarchyTreeActions.cs
namespace SubHierarchy;

/// <summary>
/// Actions that drive the built-in Hierarchy / scene visibility.
/// </summary>
public static class HierarchyTreeActions
{
	const string HiddenTag = "hidden";

	public static TreeView GetHierarchyTree()
	{
		return SceneTreeWidget.Current?.TreeView;
	}

	public static Scene GetActiveScene()
	{
		return SceneEditorSession.Active?.Scene;
	}

	/// <summary>Collapse every open group; keep the scene root expanded.</summary>
	public static void CollapseAll()
	{
		var tree = GetHierarchyTree();
		var scene = GetActiveScene();
		if ( tree is null || scene is null || !scene.IsValid() )
			return;

		tree.Close( scene, recursive: true );
		tree.Open( scene );
		tree.Update();
	}

	/// <summary>Expand every group under the scene root.</summary>
	public static void ExpandAll()
	{
		var tree = GetHierarchyTree();
		var scene = GetActiveScene();
		if ( tree is null || scene is null || !scene.IsValid() )
			return;

		tree.Open( scene, recursive: true );
		tree.Update();
	}

	/// <summary>Hide every GameObject in the scene (same <c>hidden</c> tag as H / U).</summary>
	public static void HideAll()
	{
		using var scope = SceneEditorSession.Scope();
		var session = SceneEditorSession.Active;
		if ( session?.Scene is not { } scene || !scene.IsValid() )
			return;

		var gos = scene.GetAllObjects( true )
			.Where( go => go.IsValid() && go is not Scene && !go.Tags.Has( HiddenTag ) )
			.ToArray();

		if ( gos.Length == 0 )
			return;

		using ( session.UndoScope( "Hide All Object(s)" )
			.WithGameObjectChanges( gos, GameObjectUndoFlags.All )
			.Push() )
		{
			foreach ( var go in gos )
				go.Tags.Add( HiddenTag );
		}
	}

	/// <summary>Show every GameObject that has the <c>hidden</c> tag.</summary>
	public static void ShowAll()
	{
		using var scope = SceneEditorSession.Scope();
		var session = SceneEditorSession.Active;
		if ( session?.Scene is not { } scene || !scene.IsValid() )
			return;

		var gos = scene.GetAllObjects( true )
			.Where( go => go.IsValid() && go.Tags.Has( HiddenTag ) )
			.ToArray();

		if ( gos.Length == 0 )
			return;

		using ( session.UndoScope( "Unhide All Object(s)" )
			.WithGameObjectChanges( gos, GameObjectUndoFlags.All )
			.Push() )
		{
			foreach ( var go in gos )
				go.Tags.Remove( HiddenTag );
		}
	}

	/// <summary>Hide everything except the selection, its ancestors, and its descendants.</summary>
	public static void IsolateSelection()
	{
		using var scope = SceneEditorSession.Scope();
		var session = SceneEditorSession.Active;
		if ( session?.Scene is not { } scene || !scene.IsValid() )
			return;

		var selected = EditorScene.Selection
			.OfType<GameObject>()
			.Where( go => go.IsValid() && go is not Scene )
			.ToArray();
		if ( selected.Length == 0 )
			return;

		var visible = new HashSet<GameObject>();
		foreach ( var s in selected )
		{
			// Descendants first — CollectSelfAndDescendants skips if already present.
			CollectSelfAndDescendants( s, visible );

			for ( var go = s.Parent; go is not null and not Scene; go = go.Parent )
				visible.Add( go );
		}

		var changedObjects = scene.GetAllObjects( true )
			.Where( go => go.IsValid() && go is not Scene )
			.Where( go => go.Tags.Has( HiddenTag ) == visible.Contains( go ) )
			.ToArray();
		if ( changedObjects.Length == 0 )
			return;

		using ( session.UndoScope( "Isolate Selection" )
			.WithGameObjectChanges( changedObjects, GameObjectUndoFlags.All )
			.Push() )
		{
			foreach ( var go in changedObjects )
			{
				if ( visible.Contains( go ) ) go.Tags.Remove( HiddenTag );
				else go.Tags.Add( HiddenTag );
			}
		}
	}

	/// <summary>Apply a HierarchyColor to the current selection.</summary>
	public static void ApplyColorToSelection( Color color, bool includeChildren = false )
	{
		using var scope = SceneEditorSession.Scope();
		var session = SceneEditorSession.Active;
		if ( session is null )
			return;

		var gos = CollectSelection( includeChildren );
		if ( gos.Length == 0 )
			return;

		using ( session.UndoScope( "Set Hierarchy Color" )
			.WithGameObjectChanges( gos, GameObjectUndoFlags.All )
			.Push() )
		{
			foreach ( var go in gos )
			{
				var marker = go.Components.GetOrCreate<HierarchyColor>();
				marker.Tint = color;
			}
		}

		GetHierarchyTree()?.Update();
	}

	/// <summary>Remove HierarchyColor from the current selection.</summary>
	public static void ClearColorFromSelection( bool includeChildren = false )
	{
		using var scope = SceneEditorSession.Scope();
		var session = SceneEditorSession.Active;
		if ( session is null )
			return;

		var gos = CollectSelection( includeChildren )
			.Where( go => go.Components.Get<HierarchyColor>( true ) is not null )
			.ToArray();

		if ( gos.Length == 0 )
			return;

		using ( session.UndoScope( "Clear Hierarchy Color" )
			.WithGameObjectChanges( gos, GameObjectUndoFlags.All )
			.Push() )
		{
			foreach ( var go in gos )
			{
				var marker = go.Components.Get<HierarchyColor>( true );
				marker?.Destroy();
			}
		}

		GetHierarchyTree()?.Update();
	}

	/// <summary>Remove HierarchyColor from every object in the active scene.</summary>
	public static void ClearAllColors()
	{
		using var scope = SceneEditorSession.Scope();
		var session = SceneEditorSession.Active;
		if ( session?.Scene is not { } scene || !scene.IsValid() )
			return;

		var gos = scene.GetAllObjects( true )
			.Where( go => go.IsValid() && go is not Scene && go.Components.Get<HierarchyColor>( true ) is not null )
			.ToArray();

		if ( gos.Length == 0 )
			return;

		using ( session.UndoScope( "Clear All Hierarchy Colors" )
			.WithGameObjectChanges( gos, GameObjectUndoFlags.All )
			.Push() )
		{
			foreach ( var go in gos )
			{
				var marker = go.Components.Get<HierarchyColor>( true );
				marker?.Destroy();
			}
		}

		GetHierarchyTree()?.Update();
	}

	/// <summary>Remove HierarchyColor from every object whose tint matches <paramref name="color"/>.</summary>
	public static void ClearColorsMatching( Color color, float epsilon = 0.02f )
	{
		using var scope = SceneEditorSession.Scope();
		var session = SceneEditorSession.Active;
		if ( session?.Scene is not { } scene || !scene.IsValid() )
			return;

		var gos = scene.GetAllObjects( true )
			.Where( go => go.IsValid() && go is not Scene )
			.Where( go =>
			{
				var marker = go.Components.Get<HierarchyColor>( true );
				return marker is not null && marker.IsValid() && ColorsMatch( marker.Tint, color, epsilon );
			} )
			.ToArray();

		if ( gos.Length == 0 )
			return;

		using ( session.UndoScope( "Clear Matching Hierarchy Colors" )
			.WithGameObjectChanges( gos, GameObjectUndoFlags.All )
			.Push() )
		{
			foreach ( var go in gos )
			{
				var marker = go.Components.Get<HierarchyColor>( true );
				marker?.Destroy();
			}
		}

		GetHierarchyTree()?.Update();
	}

	static bool ColorsMatch( Color a, Color b, float epsilon )
	{
		return MathF.Abs( a.r - b.r ) < epsilon
			&& MathF.Abs( a.g - b.g ) < epsilon
			&& MathF.Abs( a.b - b.b ) < epsilon;
	}

	static GameObject[] CollectSelection( bool includeChildren )
	{
		var selected = EditorScene.Selection
			.OfType<GameObject>()
			.Where( go => go.IsValid() && go is not Scene )
			.ToArray();

		if ( !includeChildren || selected.Length == 0 )
			return selected;

		var set = new HashSet<GameObject>();
		foreach ( var root in selected )
			CollectSelfAndDescendants( root, set );

		return set.ToArray();
	}

	static void CollectSelfAndDescendants( GameObject go, HashSet<GameObject> set )
	{
		if ( !go.IsValid() || go is Scene || !set.Add( go ) )
			return;

		foreach ( var child in go.Children )
			CollectSelfAndDescendants( child, set );
	}

	/// <summary>Sort every parent's children A–Z by name (recursive).</summary>
	public static void ReorganizeAlphabetically()
	{
		ReorganizeHierarchy(
			"Reorganize Hierarchy Alphabetically",
			children => children
				.OrderBy( go => go.Name ?? string.Empty, StringComparer.OrdinalIgnoreCase )
				.ThenBy( go => go.Id )
				.ToList() );
	}

	/// <summary>Sort every parent's children by HierarchyColor hue, then name (recursive).</summary>
	public static void ReorganizeByColor()
	{
		ReorganizeHierarchy(
			"Reorganize Hierarchy By Color",
			children => children
				.OrderBy( GetColorSortKey )
				.ThenBy( go => go.Name ?? string.Empty, StringComparer.OrdinalIgnoreCase )
				.ThenBy( go => go.Id )
				.ToList() );
	}

	static void ReorganizeHierarchy( string undoName, Func<List<GameObject>, List<GameObject>> sorter )
	{
		using var scope = SceneEditorSession.Scope();
		var session = SceneEditorSession.Active;
		if ( session?.Scene is not { } scene || !scene.IsValid() )
			return;

		var affected = scene.GetAllObjects( true )
			.Where( go => go.IsValid() && go is not Scene )
			.ToArray();

		if ( affected.Length == 0 )
			return;

		using ( session.UndoScope( undoName )
			.WithGameObjectChanges( affected, GameObjectUndoFlags.All )
			.Push() )
		{
			SortChildrenRecursive( scene, sorter );
		}

		GetHierarchyTree()?.Update();
	}

	static void SortChildrenRecursive( GameObject parent, Func<List<GameObject>, List<GameObject>> sorter )
	{
		if ( !parent.IsValid() )
			return;

		SortChildrenOf( parent, sorter );

		foreach ( var child in parent.Children.ToArray() )
			SortChildrenRecursive( child, sorter );
	}

	static void SortChildrenOf( GameObject parent, Func<List<GameObject>, List<GameObject>> sorter )
	{
		var children = parent.Children
			.Where( go => go.IsValid() && ShouldReorder( go ) )
			.ToList();

		if ( children.Count < 2 )
			return;

		var sorted = sorter( children );
		var alreadyOrdered = true;
		for ( var i = 0; i < sorted.Count; i++ )
		{
			if ( !ReferenceEquals( children[i], sorted[i] ) )
			{
				alreadyOrdered = false;
				break;
			}
		}

		if ( alreadyOrdered )
			return;

		for ( var i = 0; i < sorted.Count; i++ )
		{
			var desired = sorted[i];
			var live = parent.Children.Where( ShouldReorder ).ToList();
			if ( i >= live.Count )
				break;

			if ( ReferenceEquals( live[i], desired ) )
				continue;

			if ( i == 0 )
				live[0].AddSibling( desired, before: true );
			else
				sorted[i - 1].AddSibling( desired, before: false );
		}
	}

	static bool ShouldReorder( GameObject go )
	{
		if ( !go.IsValid() || go is Scene )
			return false;

		// Leave engine-generated hierarchy nodes alone.
		if ( go.Flags.Contains( GameObjectFlags.Bone ) )
			return false;
		if ( go.Flags.Contains( GameObjectFlags.Attachment ) )
			return false;
		if ( go.Flags.Contains( GameObjectFlags.PhysicsBone ) )
			return false;
		if ( go.Flags.Contains( GameObjectFlags.ProceduralBone ) )
			return false;

		return true;
	}

	/// <summary>Colored objects by hue; uncolored last.</summary>
	static float GetColorSortKey( GameObject go )
	{
		var marker = go.Components.Get<HierarchyColor>( true );
		if ( marker is null || !marker.IsValid() )
			return float.MaxValue;

		var hsv = marker.Tint.ToHsv();
		return hsv.Hue;
	}
}