Editor/ScenePicking.cs
using System;
using System.Collections;
using System.Collections.Generic;
using System.Linq;

namespace HierarchyExtension;

/// <summary>
/// Makes the scene view honour <see cref="Unselectable"/>, and makes its drag selections take only what's fully inside.
/// <list type="bullet">
/// <item><description><b>Clicks</b>: each scene viewport's gizmos select through a <see cref="ClickFilter"/>. A click
/// that would select an unselectable object is held back. On the next editor frame <see cref="ClickThrough"/> plays it
/// again without the unselectable objects, and its selection is made instead, as one undo step. With nothing behind,
/// the click deselects, like a click on empty space.</description></item>
/// <item><description><b>Drags</b> that start on an unselectable object are drags on empty space: drag selections,
/// and the camera, take them as they would there (<see cref="LetGoOfDraggedPress"/>).</description></item>
/// <item><description><b>Drag selections</b> with the Object Select tool take only objects fully inside the box, never
/// unselectable ones, and a prefab only whole (<see cref="WatchDragSelection"/>).</description></item>
/// <item><description><b>Right-clicking</b> the scene view offers to make unselectable objects selectable again.</description></item>
/// </list>
/// <para>
/// The filter goes where the engine lets it: a viewport's <c>GizmoInstance.Selection</c> is a public setting, and only
/// the gizmos' clicks select through it (tools source, <c>SceneViewportWidget</c>, 26.09.29). The Hierarchy, the
/// inspector, tools and code select through the session's selection, so they can still pick anything.
/// </para>
/// </summary>
public sealed class ScenePicking : EditorEvent.ISceneView
{
	/// <summary>How far the cursor goes before a drag selection starts (<c>EditorTool.UpdateBoxSelection</c>), in pixels.</summary>
	const float DragDistance = 5;

	static ScenePicking _listener;
	static PendingClick _pending;
	static DragWatch _drag;

	/// <summary>Whether the press going on has been checked for landing on an unselectable object.</summary>
	static bool _pressChecked;

	/// <summary>
	/// While a drag selection is going on in the scene view: what it will select when the mouse comes up. Null otherwise.
	/// The Hierarchy shows this meanwhile (<see cref="HierarchySelection"/>).
	/// </summary>
	internal static HashSet<object> DragResult => _drag?.Result;

	/// <summary>A click that was held back, with what the viewport's gizmos saw when it landed.</summary>
	sealed record PendingClick( SceneViewportWidget Viewport, SelectionSystem Selection, Gizmo.Inputs Input );

	/// <summary>A press begun in a viewport (known by its camera), and the drag selection it may become.</summary>
	sealed class DragWatch
	{
		public Scene Scene;
		public SceneCamera Camera;

		/// <summary>Editor frames in a row the left button has been up while this drag still went on.</summary>
		public int FramesUp;

		/// <summary>What was selected when the press began.</summary>
		public HashSet<object> Before;

		/// <summary>The cursor's ray when the press began, and on the last frame the button was down, as the tool takes them.</summary>
		public Ray Start, Last;

		/// <summary>While it's a drag selection: what the selection will be when the mouse comes up. Null otherwise.</summary>
		public HashSet<object> Result;

		readonly Dictionary<GameObject, BBox?> _bounds = new();

		/// <summary>
		/// <see cref="VisibleBounds"/>, worked out once per drag, as the tool does its bounds: nothing moves while the
		/// mouse is busy with a drag selection, and they're slow to get.
		/// </summary>
		public BBox? BoundsOf( GameObject go )
		{
			if ( !_bounds.TryGetValue( go, out var bounds ) ) _bounds[go] = bounds = VisibleBounds( go );
			return bounds;
		}
	}

	[EditorEvent.Frame]
	static void Frame()
	{
		if ( _listener is null )
		{
			_listener = new ScenePicking();
			EditorEvent.Register( _listener );
		}

		InstallFilters();
		ResolvePendingClick();

		// a drag whose viewport stopped drawing before its release was seen: end it once the button has been up a while
		if ( _drag is { } drag )
		{
			drag.FramesUp = Editor.Application.MouseButtons.HasFlag( MouseButtons.Left ) ? 0 : drag.FramesUp + 1;
			if ( drag.FramesUp > 3 ) Finish( drag );
		}
	}

	/// <summary>A code reload keeps the listener; register it again in case the reload dropped it.</summary>
	[EditorEvent.Hotload]
	static void OnHotload()
	{
		if ( _listener is null ) return;

		EditorEvent.Unregister( _listener );
		EditorEvent.Register( _listener );
	}

	/// <summary>
	/// Puts a filter on every viewport of the scene view in use, and on any it's lost: a viewport sets its selection
	/// again when play starts or stops. Also lets through a clear that's still waiting for its add.
	/// </summary>
	static void InstallFilters()
	{
		var view = SceneViewWidget.Current;
		if ( !view.IsValid() || view._viewports is null ) return;

		foreach ( var viewport in view._viewports.Values )
		{
			var gizmos = viewport.IsValid() ? viewport.GizmoInstance : null;
			if ( gizmos?.Selection is null ) continue;

			if ( gizmos.Selection is ClickFilter filter )
			{
				filter.Flush();
				continue;
			}

			// one from before a code reload that wasn't carried over: leave it rather than stack another on it
			if ( gizmos.Selection.GetType().FullName == typeof( ClickFilter ).FullName ) continue;

			gizmos.Selection = new ClickFilter( viewport, gizmos.Selection );
		}
	}

	static void ResolvePendingClick()
	{
		var click = _pending;
		if ( click is null ) return;
		_pending = null;

		var session = SceneEditorSession.All.FirstOrDefault( x => x.Selection == click.Selection );
		var scene = session?.Scene;
		if ( !scene.IsValid() || !click.Viewport.IsValid() ) return;

		var changes = ClickThrough.Run( scene, click.Viewport.GizmoInstance.Settings, click.Input, click.Selection );

		using var sceneScope = scene.Push();
		using ( session.UndoScope( UndoName( changes ) ).Push() )
		{
			if ( changes.Count == 0 )
			{
				click.Selection.Clear();
				return;
			}

			foreach ( var (change, target) in changes )
			{
				switch ( change )
				{
					case ClickThrough.Change.Clear: click.Selection.Clear(); break;
					case ClickThrough.Change.Add: click.Selection.Add( target ); break;
					case ClickThrough.Change.Remove: click.Selection.Remove( target ); break;
				}
			}
		}
	}

	/// <summary>Named as the stock click's undo steps are: "Select {name}", or "Deselect all" for empty space.</summary>
	static string UndoName( List<(ClickThrough.Change Change, object Target)> changes )
	{
		var last = changes.LastOrDefault( x => x.Change != ClickThrough.Change.Clear );
		if ( last.Target is GameObject go ) return last.Change == ClickThrough.Change.Add ? $"Select {go.Name}" : $"Deselect {go.Name}";

		return "Deselect all";
	}

	/// <summary>Runs in each viewport's frame after its tool, before the scene's gizmos.</summary>
	void EditorEvent.ISceneView.DrawGizmos( Scene scene )
	{
		LetGoOfDraggedPress( scene );
		WatchDragSelection( scene );

		// what's selected, or during a drag selection what it will select
		var selected = DragResult ?? (IEnumerable<object>)Unselectable.SessionOf( scene )?.Selection;
		if ( selected is not null ) SelectionOutline.Draw( selected );
	}

	/// <summary>
	/// A press on an unselectable object still counts as a press on something: the scene view hovers what's under the
	/// cursor before anyone clicks, with no way to leave objects out. While anything is pressed, a drag doesn't start a
	/// drag selection (<c>EditorTool.UpdateBoxSelection</c>) or move the camera (<c>SceneViewportWidget</c>), so a drag
	/// that began on a backdrop did nothing. Once such a press has moved as far as a drag selection waits for, let go of
	/// it (<c>Gizmo.Pressed.ClearPath</c>, as the engine does itself for a dragged object press): from the next frame the
	/// drag is a drag on empty space, and a drag selection starts from where the press began. A press that doesn't move
	/// stays a click, which the <see cref="ClickFilter"/> plays through to what's behind.
	/// </summary>
	static void LetGoOfDraggedPress( Scene scene )
	{
		if ( !Gizmo.HasMouseFocus ) return;
		if ( Gizmo.WasLeftMousePressed ) _pressChecked = false;

		if ( _pressChecked || !Gizmo.IsLeftMouseDown || !Gizmo.Pressed.Any ) return;
		if ( Gizmo.Pressed.CursorDelta.Length < DragDistance ) return;

		_pressChecked = true;
		if ( PressedUnselectable( scene ).IsValid() ) Gizmo.Pressed.ClearPath();
	}

	/// <summary>
	/// The unselectable object the press is on, or null. The gizmos name what's pressed by its path: an object's is
	/// "object-{hash}" (<c>Gizmo.ObjectScope</c>), its handle's ".../Handle", and its components' hitboxes ".../Scope0",
	/// ".../Scope1" and so on, in order (<c>GameObject.DrawGizmos</c>). <c>Gizmo.Pressed.This</c> compares the path it's
	/// asked in, so each is tried. A hitbox a component draws inside a scope of its own isn't found.
	/// </summary>
	internal static GameObject PressedUnselectable( Scene scene ) => PressedObject( scene, Unselectable.Has );

	/// <summary>The object the press is on, among those <paramref name="among"/> accepts, or null. See <see cref="PressedUnselectable"/>.</summary>
	internal static GameObject PressedObject( Scene scene, System.Func<GameObject, bool> among )
	{
		using var scope = Gizmo.Scope( "hierarchyextension" );

		foreach ( var go in scene.GetAllObjects( true ) )
		{
			if ( !among( go ) ) continue;

			var path = $"object-{go.GetHashCode()}";
			if ( IsPressed( path ) || IsPressed( $"{path}/Handle" ) ) return go;

			for ( var i = 0; i < go.Components.Count; i++ )
			{
				if ( IsPressed( $"{path}/Scope{i}" ) ) return go;
			}
		}

		return null;
	}

	static bool IsPressed( string path )
	{
		Gizmo.Path = path;
		return Gizmo.Pressed.This;
	}

	/// <summary>
	/// Makes the Object Select tool's drag selections take only what's fully inside the box. The stock one takes anything
	/// whose bounds the box touches, and an object's bounds include everything under it (<c>ObjectEditorTool.OnBoxSelect</c>,
	/// <c>GameObject.GetBounds</c>), so a small box over a few decorations also took their group and the map they're in.
	/// It has no way in, so it runs as before and the result is redone here, from the same box:
	/// <list type="bullet">
	/// <item><description>An object is taken when it's fully inside with everything under it, and then stands for all
	/// of that: what's under it isn't taken as well.</description></item>
	/// <item><description>Inside a prefab, or anything else a click selects as a whole (<see cref="IsSelectionBase"/>),
	/// only that whole counts.</description></item>
	/// <item><description>Hidden and unselectable objects are never taken.</description></item>
	/// <item><description>Ctrl takes them out of the selection, Shift adds them, and otherwise they're the selection,
	/// as with the stock one (<c>EditorTool.ApplyDragSelection</c>).</description></item>
	/// </list>
	/// While the drag goes on the scene view (<see cref="ClickFilter.Contains"/>) and the Hierarchy
	/// (<see cref="HierarchySelection"/>) show the result; it goes into the selection when the mouse comes up, the frame
	/// the stock one finishes. Putting it in every frame would have the tool undo it every frame, and the inspector rebuild
	/// and the Hierarchy scroll each time, so the inspector and the move handles show the stock one's picks until then.
	/// </summary>
	static void WatchDragSelection( Scene scene )
	{
		var camera = Gizmo.Camera;

		if ( Gizmo.WasLeftMousePressed && Gizmo.HasMouseFocus )
		{
			var before = Unselectable.SessionOf( scene )?.Selection;
			_drag = before is null ? null : new DragWatch { Scene = scene, Camera = camera, Before = before.ToHashSet(), Start = Gizmo.CurrentRay, Last = Gizmo.CurrentRay };
			return;
		}

		var drag = _drag;
		if ( drag is null || drag.Camera != camera ) return;

		// the frame the mouse comes up still counts: the tool finishes its drag selection on it, from the last box. With the
		// cursor out of the view the tool leaves its box as it was, and so does this
		if ( Gizmo.HasMouseFocus && (Gizmo.IsLeftMouseDown || Gizmo.WasLeftMouseReleased) )
		{
			if ( Gizmo.IsLeftMouseDown ) drag.Last = Gizmo.CurrentRay;
			drag.Result = IsBoxSelecting( drag, out var frustum ) ? BoxResult( scene, drag, frustum ) : null;

			// the Hierarchy shows the result as it changes, not what the stock one has put in the selection meanwhile
			HierarchySelection.Current?.Sync();
		}

		// the release, or a later frame if the release went by unseen (a code reload, say)
		if ( !Gizmo.IsLeftMouseDown ) Finish( drag );
	}

	/// <summary>Ends the drag: a drag selection's result goes into the selection, and the Hierarchy shows the selection again.</summary>
	static void Finish( DragWatch drag )
	{
		if ( _drag == drag ) _drag = null;

		var result = drag.Result;
		var selection = Unselectable.SessionOf( drag.Scene )?.Selection;
		if ( result is not null && selection is not null )
		{
			foreach ( var obj in selection.ToArray() )
			{
				if ( !result.Contains( obj ) ) selection.Remove( obj );
			}

			foreach ( var obj in result )
			{
				if ( !selection.Contains( obj ) ) selection.Add( obj );
			}
		}

		HierarchySelection.Current?.Sync();
	}

	/// <summary>
	/// Whether the Object Select tool is making a drag selection this frame, decided as it decides it
	/// (<c>EditorTool.UpdateBoxSelection</c>): the cursor 5 px or more from where the press began, and nothing pressed.
	/// Gives the box's frustum.
	/// </summary>
	static bool IsBoxSelecting( DragWatch drag, out Frustum frustum )
	{
		frustum = default;

		if ( EditorToolManager.CurrentModeName != nameof( ObjectEditorTool ) ) return false;
		if ( Gizmo.Pressed.Any ) return false;

		var c1 = Gizmo.Camera.ToScreen( drag.Start.Project( 100 ) );
		var c2 = Gizmo.Camera.ToScreen( drag.Last.Project( 100 ) );
		if ( Vector2.Distance( c1, c2 ) < DragDistance ) return false;

		frustum = Gizmo.Camera.GetFrustum( Rect.FromPoints( c1, c2 ) );
		return true;
	}

	/// <summary>What the selection will be when the drag selection in <paramref name="frustum"/> finishes.</summary>
	static HashSet<object> BoxResult( Scene scene, DragWatch drag, Frustum frustum )
	{
		var taken = Take( scene, frustum, go => drag.BoundsOf( go ) );
		var result = new HashSet<object>( drag.Before );

		if ( Gizmo.IsCtrlPressed )
		{
			result.ExceptWith( taken );
			return result;
		}

		if ( !Gizmo.IsShiftPressed ) result.RemoveWhere( x => x is GameObject );

		result.UnionWith( taken );
		return result;
	}

	/// <summary>
	/// The objects a drag selection in <paramref name="frustum"/> takes, before Ctrl or Shift: those fully inside with
	/// everything under them (<paramref name="boundsOf"/>, usually <see cref="VisibleBounds"/>), none that are part of
	/// one taken, hidden or unselectable, and a prefab only whole.
	/// </summary>
	internal static HashSet<GameObject> Take( Scene scene, Frustum frustum, System.Func<GameObject, BBox?> boundsOf )
	{
		var taken = new HashSet<GameObject>();

		foreach ( var go in scene.GetAllObjects( true ) )
		{
			if ( go is Scene || go.Tags.Has( "hidden" ) || Unselectable.Has( go ) ) continue;

			// part of a prefab: the prefab is taken whole or not at all
			if ( InsideSelectionBase( go ) ) continue;

			if ( boundsOf( go ) is { } bounds && frustum.IsInside( bounds, false ) ) taken.Add( go );
		}

		// an object taken with everything under it stands for all of it
		taken.RemoveWhere( go => HasAncestorIn( go, taken ) );
		return taken;
	}

	/// <summary>
	/// What the scene view shows of <paramref name="go"/> and everything under it, as one world box, or null when there's
	/// nothing. That's what the stock drag selection measures (<c>GameObject.GetBounds</c>: colliders, models, particles),
	/// plus sprites, which it leaves out, and the icon of each object that has one, as a point. A
	/// <c>SpriteRenderer</c> has no bounds of its own (it isn't a <c>Component.IHasBounds</c>, engine source 26.09.29),
	/// so with the stock drag selection no sprite can be drag-selected, and a 2D level's decorations never are.
	/// </summary>
	internal static BBox? VisibleBounds( GameObject go )
	{
		BBox? result = null;
		void Add( BBox bounds ) => result = result is { } r ? r.AddBBox( bounds ) : bounds;

		foreach ( var o in go.GetAllObjects( true ) )
		{
			if ( o.Tags.Has( "hidden" ) ) continue;

			foreach ( var c in o.Components.GetAll( FindMode.EnabledInSelf ) )
			{
				if ( c is Component.IHasBounds hasBounds ) Add( hasBounds.LocalBounds.Transform( c.WorldTransform ) );
				if ( c is SpriteRenderer sprite ) Add( SpriteBounds( sprite ) );
			}

			if ( o.HasGizmoHandle ) Add( BBox.FromPositionAndSize( o.WorldPosition ) );
		}

		return result;
	}

	/// <summary>A sprite's box in the world, worked out as its own hitbox is (<c>SpriteRenderer.DrawGizmos</c>).</summary>
	internal static BBox SpriteBounds( SpriteRenderer sprite )
	{
		var world = sprite.WorldTransform;
		var billboard = sprite.Billboard is SpriteRenderer.BillboardMode.Always or SpriteRenderer.BillboardMode.YOnly;

		var transform = billboard ? world.WithRotation( new Rotation() ) : world;
		transform = transform.WithScale( new Vector3( world.Scale.x, world.Scale.x, world.Scale.z ) );

		var pivot = ((sprite.CurrentAnimation?.Origin ?? new Vector2( 0.5f, 0.5f )) - 0.5f) * sprite.Size;
		var size = new Vector2( sprite.Size.x, sprite.Size.y );

		if ( billboard )
		{
			size += Vector2.Abs( pivot * 2 );

			var angle = MathX.DegreeToRadian( world.Rotation.Roll() );
			var cos = MathF.Abs( MathF.Cos( angle ) );
			var sin = MathF.Abs( MathF.Sin( angle ) );
			size = new Vector2( size.x * cos + size.y * sin, size.x * sin + size.y * cos );
		}

		var local = BBox.FromPositionAndSize(
			billboard ? Vector3.Zero : new Vector3( 0, pivot.x, pivot.y ),
			new Vector3( billboard ? size.x : 0.5f, size.x, size.y ) );

		return local.Transform( transform );
	}

	/// <summary>
	/// Whether a click selects <paramref name="go"/> as a whole when it lands on something under it, as the scene view
	/// decides it (<c>GameObject.FindSelectionBase</c>): a network root, the outermost root of a prefab instance, or an
	/// object with a <see cref="SelectionBaseAttribute"/> component.
	/// </summary>
	static bool IsSelectionBase( GameObject go )
	{
		// GameObject.IsOutermostPrefabInstanceRoot isn't public; the stock Hierarchy asks this
		if ( go.IsNetworkRoot || EditorUtility.Prefabs.IsOuterMostPrefabRoot( go ) ) return true;

		return go.Components.GetAll().Any( x => EditorTypeLibrary.GetType( x.GetType() )?.HasAttribute<SelectionBaseAttribute>() == true );
	}

	static bool InsideSelectionBase( GameObject go )
	{
		for ( var o = go.Parent; o.IsValid() && o is not Scene; o = o.Parent )
		{
			if ( IsSelectionBase( o ) ) return true;
		}

		return false;
	}

	static bool HasAncestorIn( GameObject go, HashSet<GameObject> objects )
	{
		for ( var o = go.Parent; o.IsValid() && o is not Scene; o = o.Parent )
		{
			if ( objects.Contains( o ) ) return true;
		}

		return false;
	}

	void EditorEvent.ISceneView.ShowContextMenu( EditorEvent.ShowContextMenuEvent ev )
	{
		var scene = ev.Session?.Scene;
		if ( !Unselectable.AnyIn( scene ) ) return;

		ev.Menu.AddSeparator();

		if ( ev.Trace is { Hit: true } trace && Unselectable.Source( trace.GameObject ) is { } source )
		{
			ev.Menu.AddOption( $"Make {source.Name} Selectable", "touch_app", () => Unselectable.Set( new[] { source }, false ) );
		}

		ev.Menu.AddOption( "Make All Selectable", "touch_app", Unselectable.MakeAllSelectable );
	}

	/// <summary>
	/// A scene viewport's gizmo selection. It passes everything on to the scene's own selection, except a click on an
	/// unselectable object, which it holds back for <see cref="ResolvePendingClick"/>.
	/// </summary>
	sealed class ClickFilter : SelectionSystem, IEnumerable<object>
	{
		readonly SceneViewportWidget _viewport;
		readonly SelectionSystem _selection;

		/// <summary>
		/// A click clears the selection just before it adds what it clicked (<c>Gizmo.Select</c>). The clear waits for
		/// the add: if the click is held back, the clear isn't wanted either.
		/// </summary>
		bool _clearing;

		public ClickFilter( SceneViewportWidget viewport, SelectionSystem selection )
		{
			_viewport = viewport;
			_selection = selection;
		}

		public override void Clear() => _clearing = true;

		public override bool Add( object obj )
		{
			if ( HoldBack( obj ) ) return false;

			Flush();
			return _selection.Add( obj );
		}

		public override bool Remove( object obj )
		{
			if ( HoldBack( obj ) ) return false;

			Flush();
			return _selection.Remove( obj );
		}

		public override bool Set( object obj )
		{
			Clear();
			return Add( obj );
		}

		/// <summary>
		/// Also what the gizmos draw as selected. While a drag selection goes on, that's what it will select when the mouse
		/// comes up (<see cref="WatchDragSelection"/>), not what the stock one has put in the selection meanwhile.
		/// </summary>
		public override bool Contains( object obj ) => _drag?.Result is { } result ? result.Contains( obj ) : _selection.Contains( obj );

		public override bool Any() => _drag?.Result is { } result ? result.Count > 0 : _selection.Any();

		public override int GetHashCode() => _selection.GetHashCode();

		IEnumerator<object> IEnumerable<object>.GetEnumerator() => _selection.GetEnumerator();

		IEnumerator IEnumerable.GetEnumerator() => _selection.GetEnumerator();

		/// <summary>A clear still waiting for its add goes through.</summary>
		public void Flush()
		{
			if ( !_clearing ) return;

			_clearing = false;
			_selection.Clear();
		}

		bool HoldBack( object obj )
		{
			if ( obj is not GameObject go || !Unselectable.Has( go ) ) return false;

			_clearing = false;
			_pending = new PendingClick( _viewport, _selection, _viewport.GizmoInstance.Input );
			return true;
		}
	}
}