Editor/HierarchyColorPainter.cs
using System.Collections;
using System.Reflection;

namespace SubHierarchy;

/// <summary>
/// Tints Hierarchy rows for GameObjects that have a <see cref="HierarchyColor"/> component.
/// Hooks the built-in Hierarchy TreeView paint (GameObjectNode is internal).
/// </summary>
public static class HierarchyColorPainter
{
	const float TintAlpha = 0.38f;
	const BindingFlags LayoutFlags = BindingFlags.Instance | BindingFlags.Public | BindingFlags.NonPublic;

	static TreeView _hookedTree;
	static PropertyInfo _itemLayoutsProperty;
	static FieldInfo _itemLayoutsField;
	static bool _layoutsLookupDone;

	[Event( "editor.created" )]
	static void OnEditorCreated( object _ ) => EnsureHooked();

	[EditorEvent.Frame]
	static void OnFrame() => EnsureHooked();

	public static void EnsureHooked()
	{
		var tree = SceneTreeWidget.Current?.TreeView;
		if ( tree is null || !tree.IsValid() )
		{
			_hookedTree = null;
			return;
		}

		// Always re-apply: Hierarchy BuildUI can replace TreeView or reset OnPaintOverride.
		if ( !ReferenceEquals( _hookedTree, tree ) )
			_hookedTree = tree;

		tree.OnPaintOverride = PaintHierarchyBackground;
	}

	static bool PaintHierarchyBackground()
	{
		var tree = _hookedTree;
		if ( tree is null || !tree.IsValid() )
			return false;

		Paint.ClearPen();
		Paint.SetBrush( Theme.ControlBackground );
		Paint.DrawRect( tree.LocalRect, Theme.ControlRadius );

		PaintColoredRows( tree );
		return false;
	}

	static void PaintColoredRows( TreeView tree )
	{
		var layouts = GetItemLayouts( tree );
		if ( layouts is not null )
		{
			foreach ( var entry in layouts )
			{
				if ( entry is not VirtualWidget item )
					continue;

				if ( !TryResolveGameObject( item.Object, out var go ) )
					continue;

				var marker = go.Components.Get<HierarchyColor>( true );
				if ( marker is null || !marker.IsValid() )
					continue;

				// ItemLayouts rects are already in paint / scrolled local space.
				PaintRow( tree, item.Rect, marker.Tint );
			}

			return;
		}

		// Fallback: TryGetItemRect is content-space — convert using scrollbar value.
		PaintColoredRowsFallback( tree );
	}

	static void PaintColoredRowsFallback( TreeView tree )
	{
		var scene = SceneEditorSession.Active?.Scene;
		if ( scene is null || !scene.IsValid() )
			return;

		var scroll = tree.VerticalScrollbar?.Value ?? 0f;
		var originY = tree.LocalRect.Top + scroll - tree.Margin.Top;

		foreach ( var go in scene.GetAllObjects( true ) )
		{
			if ( !go.IsValid() || go is Scene )
				continue;

			var marker = go.Components.Get<HierarchyColor>( true );
			if ( marker is null || !marker.IsValid() )
				continue;

			if ( !tree.TryGetItemRect( go, out var contentRect ) )
				continue;

			var rect = contentRect;
			rect.Top -= originY;
			PaintRow( tree, rect, marker.Tint );
		}
	}

	static void PaintRow( TreeView tree, Rect rect, Color color )
	{
		rect.Left = 0;
		rect.Right = tree.Width;

		Paint.ClearPen();
		Paint.SetBrush( color.WithAlpha( TintAlpha ) );
		Paint.DrawRect( rect );
	}

	static bool TryResolveGameObject( object obj, out GameObject go )
	{
		go = null;

		if ( obj is TreeNode node )
			go = node.Value as GameObject;
		else
			go = obj as GameObject;

		return go.IsValid() && go is not Scene;
	}

	/// <summary>
	/// Visible row layouts live on BaseItemWidget as a protected member.
	/// Those rects match what Hierarchy paints; TryGetItemRect is content-space (scroll Y) and misaligns.
	/// </summary>
	static IEnumerable GetItemLayouts( TreeView tree )
	{
		if ( !_layoutsLookupDone )
		{
			_layoutsLookupDone = true;
			var type = typeof( BaseItemWidget );
			_itemLayoutsProperty = type.GetProperty( "ItemLayouts", LayoutFlags );
			_itemLayoutsField = type.GetField( "ItemLayouts", LayoutFlags );
		}

		object value = null;
		if ( _itemLayoutsProperty is not null )
			value = _itemLayoutsProperty.GetValue( tree );
		else if ( _itemLayoutsField is not null )
			value = _itemLayoutsField.GetValue( tree );

		return value as IEnumerable;
	}
}