Editor/HierarchyColoursDock.cs
using System;
using System.IO;
using System.Text.Json;
using Editor;
using Sandbox;
using System.Collections.Generic;
using System.Linq;

namespace HierarchyChanges;


/// <summary>Prototype palette for editor-only annotations on the stock hierarchy.</summary>
[Dock( "Editor", "Hierarchy Changes", "palette", DockArea.Right )]
public sealed class HierarchyColoursDock : Widget
{
	private readonly Label status;

	public HierarchyColoursDock( Widget parent ) : base( parent )
	{
		Layout = Layout.Column();
		Layout.Margin = 8;
		Layout.Spacing = 6;
		Layout.Add( new Label( "Select objects in the hierarchy, then choose a colour.", this ) );
		Layout.Add( new Label( "Middle-click a hierarchy row to expand or collapse it.", this ) );
		for ( var i = 0; i < HierarchyColourPreview.Colours.Length; i++ )
		{
			var index = i;
			var button = Layout.Add( new Button( HierarchyColourPreview.Names[i], this ) );
			button.SetStyles( $"background-color: {HierarchyColourPreview.Css[i]}; color: white;" );
			button.Clicked += () => HierarchyColourPreview.Assign( index );
		}
		Layout.Add( new Button( "Clear selected colours", this ) ).Clicked += () => HierarchyColourPreview.Assign( -1 );
		Layout.Add( new Button( "Make selected names bold", this ) ).Clicked += () => HierarchyColourPreview.SetBold( true );
		Layout.Add( new Button( "Use regular names", this ) ).Clicked += () => HierarchyColourPreview.SetBold( false );
		Layout.Add( new Button( "Toggle colour preview", this ) ).Clicked += HierarchyColourPreview.Toggle;
		status = Layout.Add( new Label( "", this ) );
		Layout.AddStretchCell();
	}

	[EditorEvent.Frame]
	private void RefreshStatus()
	{
		status.Text = HierarchyColourPreview.Status;
	}
}

/// <summary>
/// Adds colour underlays and dispatches bold rows while retaining stock interaction.
/// This remains independent of game code so it can move into a library's Editor folder.
/// </summary>
internal static class HierarchyColourPreview
{
	internal static readonly string[] Names = ["Red", "Orange", "Yellow", "Green", "Blue", "Purple"];
	internal static readonly string[] Css = ["#873b42", "#875331", "#756725", "#346849", "#365f87", "#694482"];
	internal static readonly Color[] Colours =
	[
		new( 0.85f, 0.25f, 0.30f ), new( 0.95f, 0.48f, 0.18f ), new( 0.90f, 0.78f, 0.20f ),
		new( 0.25f, 0.75f, 0.40f ), new( 0.25f, 0.55f, 0.95f ), new( 0.68f, 0.35f, 0.90f )
	];
	private static Dictionary<string, int> assignments = new();
	// Preserve the original 0..5 colour values. Bits 0..2 hold colour (7 = none);
	// bit 3 stores bold independently, so clearing colour never clears emphasis.
	private const int NoColour = 7;
	private const int BoldFlag = 8;
	private static TreeView tree;
	private static TreeView inputTree;
	private static Func<bool> previousPaint;
	private static bool enabled = true;
	private static string settingsPath;
	private static bool storageAvailable;
	private static string error;
	internal static string Status => error ?? (enabled ? "Preview on · colours and bold save automatically." : "Preview off · saved styles retained.");

	// The hook also runs when the palette is hidden. It reconnects when the hierarchy is rebuilt.
	[EditorEvent.Frame]
	private static void Frame()
	{
		LoadProject();
		var current = SceneTreeWidget.Current?.TreeView;
		AttachMiddleClick( current is not null && current.IsValid() ? current : null );
		if ( !enabled || current is null || !current.IsValid() )
		{
			Detach();
			return;
		}
		if ( ReferenceEquals( tree, current ) ) return;
		Detach();
		tree = current;
		previousPaint = tree.OnPaintOverride;
		tree.OnPaintOverride = PaintUnderlay;
		tree.Update();
	}

	[EditorEvent.Hotload]
	private static void Hotload()
	{
		Detach();
		AttachMiddleClick( null );
	}

	private static void AttachMiddleClick( TreeView current )
	{
		if ( ReferenceEquals( inputTree, current ) ) return;
		// Input stays active with colour preview off; remove only our own subscription.
		if ( inputTree is not null && inputTree.IsValid() ) inputTree.MouseMiddlePress -= MiddleClick;
		inputTree = current;
		if ( inputTree is not null ) inputTree.MouseMiddlePress += MiddleClick;
	}

	private static void MiddleClick()
	{
		if ( inputTree is null || !inputTree.IsValid() ) return;
		var item = inputTree.GetItemAt( inputTree.FromScreen( Editor.Application.CursorPosition ) );
		if ( item is null || !item.HasChildren || item.Object is not TreeNode node
			|| node.ExpanderHidden || node.Value is not GameObject go || !go.IsValid() ) return;
		// Toggle only this branch, preserving descendant expansion state and selection.
		inputTree.Toggle( go );
		inputTree.Update();
	}

	private static void Detach()
	{
		// Restore only our own callback; don't erase another extension's later replacement.
		if ( tree is not null && tree.IsValid() && tree.OnPaintOverride == PaintUnderlay )
		{
			tree.OnPaintOverride = previousPaint;
			tree.Update();
		}
		tree = null;
		previousPaint = null;
	}

	internal static void Toggle()
	{
		enabled = !enabled;
		Frame();
	}

	private static string Key( GameObject go ) => $"{go.Scene.Id:N}/{go.Id:N}";

	internal static void Assign( int colour )
		=> EditSelected( value => (value & BoldFlag) | (colour < 0 ? NoColour : colour) );

	internal static void SetBold( bool bold )
		=> EditSelected( value => bold ? value | BoldFlag : value & ~BoldFlag );

	private static void EditSelected( Func<int, int> edit )
	{
		LoadProject();
		if ( !storageAvailable ) return;
		var scene = SceneEditorSession.Active?.Scene;
		if ( scene is null || !scene.IsEditor ) return;
		var selected = SceneEditorSession.Active.Selection.OfType<GameObject>().Where( go => go.IsValid() && go.Scene == scene ).ToArray();
		if ( selected.Length == 0 ) return;
		var next = new Dictionary<string, int>( assignments );
		foreach ( var go in selected )
		{
			var key = Key( go );
			var value = edit( next.GetValueOrDefault( key, NoColour ) );
			if ( value == NoColour ) next.Remove( key );
			else next[key] = value;
		}
		try
		{
			Directory.CreateDirectory( Path.GetDirectoryName( settingsPath ) );
			var temporary = settingsPath + ".tmp";
			File.WriteAllText( temporary, JsonSerializer.Serialize( next, new JsonSerializerOptions { WriteIndented = true } ) );
			File.Move( temporary, settingsPath, true );
			assignments = next;
			error = null;
			tree?.Update();
		}
		catch ( Exception e ) { error = $"Could not save colours: {e.Message}"; }
	}

	private static void LoadProject()
	{
		var root = Project.Current?.RootDirectory?.FullName;
		var path = root is null ? null : Path.Combine( root, "Settings", "hierarchy-colours.json" );
		if ( settingsPath == path ) return;
		settingsPath = path;
		assignments = new();
		storageAvailable = path is not null;
		error = null;
		try
		{
			if ( path is not null && File.Exists( path ) )
				assignments = JsonSerializer.Deserialize<Dictionary<string, int>>( File.ReadAllText( path ) ) ?? new();
		}
		catch ( Exception e )
		{
			// Never overwrite unreadable settings with an empty palette.
			storageAvailable = false;
			error = $"Could not load colours: {e.Message}";
		}
	}

	private static bool PaintUnderlay()
	{
		if ( previousPaint?.Invoke() == true ) return true;
		var scene = SceneEditorSession.Active?.Scene;
		if ( scene is null || !scene.IsEditor || assignments.Count == 0 ) return false;
		// ItemLayouts is protected. Public hit-testing discovers only visible rows, then
		// jumps to each row's bottom; it does not traverse every object in a large scene.
		var x = Math.Max( tree.Margin.Left + 1, tree.Width * 0.5f );
		var rows = new List<VirtualWidget>();
		var hasBold = false;
		for ( float y = 0; y < tree.Height; )
		{
			var item = tree.GetItemAt( new Vector2( x, y ) );
			if ( item is null ) { y += 1; continue; }
			y = Math.Max( y + 1, item.Rect.Bottom + 0.5f );
			rows.Add( item );
			if ( item.Object is not TreeNode node || node.Value is not GameObject go || !go.IsValid() ) continue;
			if ( !assignments.TryGetValue( Key( go ), out var style ) ) continue;
			hasBold |= (style & BoldFlag) != 0;
			var index = style & NoColour;
			if ( index >= Colours.Length ) continue;
			var rect = item.Rect;
			rect.Left = 0;
			rect.Right = tree.Width;
			Paint.ClearPen();
			Paint.SetBrush( Colours[index].WithAlpha( go.Active ? 0.24f : 0.10f ) );
			Paint.DrawRect( rect );
		}
		if ( !hasBold ) return false;
		// Stock GameObjectNode hardcodes font weight. When bold is visible, dispatch
		// rows ourselves, preserving the actual nodes for selection, drag and rename.
		Paint.Antialiasing = true;
		Paint.TextAntialiasing = true;
		foreach ( var item in rows ) PaintRow( item );
		return true;
	}

	private static void PaintRow( VirtualWidget item )
	{
		if ( item.Object is not TreeNode node ) return;
		item.Selected = tree.IsSelected( item.Object );
		Paint.SetFlags( item.Selected, item.Hovered, item.Pressed, false, true );
		var rect = item.Rect;
		var childrenRect = item.ChildrenRect;
		var oldIndent = item.Indent;
		var indent = tree.IndentWidth * item.Column + tree.ExpandWidth;
		try
		{
			item.Indent = indent;
			item.Rect.Left += indent;
			item.ChildrenRect.Left += indent + tree.IndentWidth;
			// Specialized prefab/scene root nodes keep their own painter.
			if ( node.GetType().Name == "GameObjectNode" && node.Value is GameObject go
				&& go.IsValid() && assignments.TryGetValue( Key( go ), out var style ) && (style & BoldFlag) != 0 )
				BoldHierarchyRow.PaintBold( item, go, tree );
			else node.OnPaint( item );
		}
		finally
		{
			// Hit-testing shares these layouts; never leave paint indentation behind.
			item.Rect = rect;
			item.ChildrenRect = childrenRect;
			item.Indent = oldIndent;
		}
		if ( !item.HasChildren || node.ExpanderHidden ) return;
		var expander = rect;
		expander.Left += indent - tree.ExpandWidth;
		expander.Width = tree.ExpandWidth;
		Paint.SetPen( Theme.Text.WithAlpha( item.IsOpen ? 1 : 0.6f ) );
		Paint.DrawIcon( expander, item.IsOpen ? "arrow_drop_down" : "arrow_right", 26, TextFlag.Center );
	}
}