namespace SubHierarchy;

/// <summary>
/// Separate dock with Hierarchy helpers: collapse/expand, hide/show, row colors.
/// </summary>
[Dock( "Editor", "Sub Hierarchy", "account_tree" )]
public sealed class SubHierarchyDock : Widget
{
	public const string DockTitle = "Sub Hierarchy";
	const string CustomColorsCookie = "SubHierarchy.CustomColors";
	const int MaxCustomColors = 24;
	const float ColorMatchEpsilon = 0.02f;

	static readonly Color[] Presets =
	[
		new Color( 0.85f, 0.25f, 0.25f ), // red
		new Color( 0.90f, 0.50f, 0.15f ), // orange
		new Color( 0.90f, 0.80f, 0.20f ), // yellow
		new Color( 0.55f, 0.85f, 0.20f ), // lime
		new Color( 0.30f, 0.75f, 0.35f ), // green
		new Color( 0.15f, 0.70f, 0.65f ), // teal
		new Color( 0.20f, 0.75f, 0.90f ), // cyan
		new Color( 0.25f, 0.55f, 0.95f ), // blue
		new Color( 0.40f, 0.35f, 0.90f ), // indigo
		new Color( 0.65f, 0.35f, 0.90f ), // purple
		new Color( 0.90f, 0.35f, 0.70f ), // magenta
		new Color( 0.90f, 0.45f, 0.55f ), // pink
		new Color( 0.70f, 0.50f, 0.30f ), // brown
	];

	Color _draftColor = new Color( 0.25f, 0.55f, 0.95f );
	ColorSwatchButton _swatch;
	bool _colorChildrenRecursively;
	Layout _content;
	ColorSwatchPalette _palette;
	List<Color> _customColors = new();

	public SubHierarchyDock( Widget parent ) : base( parent )
	{
		Layout = Layout.Column();

		var scroller = Layout.Add( new ScrollArea( this ), 1 );
		scroller.Canvas = new Widget( scroller )
		{
			HorizontalSizeMode = SizeMode.Flexible,
			VerticalSizeMode = SizeMode.CanGrow
		};
		scroller.Canvas.Layout = Layout.Column();
		scroller.Canvas.Layout.Margin = 8;
		scroller.Canvas.Layout.Spacing = 8;
		_content = scroller.Canvas.Layout;

		LoadCustomColors();
		BuildUi();
		HierarchyColorPainter.EnsureHooked();
	}

	public static void Open()
	{
		EditorWindow.DockManager.SetDockState( DockTitle, true );
		EditorWindow.DockManager.RaiseDock( DockTitle );
	}

	void BuildUi()
	{
		_content.Clear( true );

		AddSectionLabel( "Tree" );
		AddButtonRow(
			( "Collapse All", "unfold_less", HierarchyTreeActions.CollapseAll ),
			( "Expand All", "unfold_more", HierarchyTreeActions.ExpandAll )
		);

		AddSectionLabel( "Visibility" );
		AddButtonRow(
			( "Hide All", "visibility_off", HierarchyTreeActions.HideAll ),
			( "Show All", "visibility", HierarchyTreeActions.ShowAll )
		);
		AddButtonRow(
			( "Isolate Selection", "filter_center_focus", HierarchyTreeActions.IsolateSelection )
		);

		AddSectionLabel( "Hierarchy Color" );
		AddColorControls();
		RebuildColorSwatches();

		AddSectionLabel( "Reorganize" );
		AddButtonRow(
			( "Sort Alphabetically", "sort_by_alpha", HierarchyTreeActions.ReorganizeAlphabetically ),
			( "Sort by Color", "palette", HierarchyTreeActions.ReorganizeByColor )
		);

		_content.AddStretchCell( 1 );
	}

	void AddSectionLabel( string text )
	{
		var label = _content.Add( new Label( text ) );
		label.Color = Theme.TextControl.WithAlpha( 0.85f );
	}

	void AddButtonRow( params (string Title, string Icon, Action Action)[] buttons )
	{
		var row = _content.AddRow();
		row.Spacing = 6;

		foreach ( var button in buttons )
		{
			var btn = row.Add( new Button( button.Title, button.Icon ), 1 );
			btn.Clicked = button.Action;
		}
	}

	void AddColorControls()
	{
		var row = _content.AddRow();
		row.Spacing = 6;
		row.Alignment = TextFlag.LeftCenter;

		_swatch = row.Add( new ColorSwatchButton( () => _draftColor, OnCustomColorPicked )
		{
			RememberOnClose = true,
			OnRemember = RememberCustomColor
		} );
		_swatch.ToolTip = "Pick a color and apply it to the selection";
		_swatch.FixedWidth = Theme.RowHeight * 2;
		_swatch.FixedHeight = Theme.RowHeight;

		var clear = row.Add( new Button( "Clear", "format_color_reset" ), 1 );
		clear.ToolTip = "Remove Hierarchy Color from selection";
		clear.Clicked = () => HierarchyTreeActions.ClearColorFromSelection( _colorChildrenRecursively );

		var clearAll = row.Add( new Button( "Clear All", "delete_sweep" ), 1 );
		clearAll.ToolTip = "Remove Hierarchy Color from every object in the scene";
		clearAll.Clicked = HierarchyTreeActions.ClearAllColors;

		var recursiveRow = _content.AddRow();
		recursiveRow.Spacing = 6;
		recursiveRow.Alignment = TextFlag.LeftCenter;

		var recursive = recursiveRow.Add( new Checkbox( "Color children recursively" ) );
		recursive.State = _colorChildrenRecursively ? CheckState.On : CheckState.Off;
		recursive.ToolTip = "When applying or clearing, include all descendants of the selection";
		recursive.Clicked = () => _colorChildrenRecursively = recursive.State == CheckState.On;

		_palette = _content.Add( new ColorSwatchPalette( OnCustomColorPicked, OnColorSwatchRemoved ) );
		_palette.HorizontalSizeMode = SizeMode.Flexible;
	}

	void RebuildColorSwatches()
	{
		if ( _palette is null || !_palette.IsValid() )
			return;

		var entries = new List<(Color Color, string ToolTip, bool Removable)>();
		foreach ( var preset in Presets )
			entries.Add( (preset, "Left-click: apply", false) );

		foreach ( var custom in _customColors )
			entries.Add( (custom, "Left-click: apply. Right-click: remove from palette and all objects.", true) );

		_palette.SetColors( entries );
	}

	void OnCustomColorPicked( Color color )
	{
		_draftColor = color.WithAlpha( 1f );
		_swatch?.Update();
		HierarchyTreeActions.ApplyColorToSelection( _draftColor, _colorChildrenRecursively );
	}

	void OnColorSwatchRemoved( Color color )
	{
		color = color.WithAlpha( 1f );

		// Built-in presets stay in the palette permanently.
		if ( Presets.Any( p => ColorsMatch( p, color ) ) )
			return;

		var removedCustom = _customColors.RemoveAll( c => ColorsMatch( c, color ) ) > 0;
		if ( !removedCustom )
			return;

		HierarchyTreeActions.ClearColorsMatching( color, ColorMatchEpsilon );
		SaveCustomColors();

		if ( ColorsMatch( _draftColor, color ) )
			_draftColor = Presets[0];

		_swatch?.Update();
		RebuildColorSwatches();
	}

	void RememberCustomColor( Color color )
	{
		color = color.WithAlpha( 1f );

		if ( Presets.Any( p => ColorsMatch( p, color ) ) )
			return;

		if ( _customColors.Any( c => ColorsMatch( c, color ) ) )
			return;

		_customColors.Add( color );
		while ( _customColors.Count > MaxCustomColors )
			_customColors.RemoveAt( 0 );

		SaveCustomColors();
		RebuildColorSwatches();
	}

	void LoadCustomColors()
	{
		_customColors = EditorCookie.Get( CustomColorsCookie, new List<Color>() ) ?? new List<Color>();
	}

	void SaveCustomColors()
	{
		EditorCookie.Set( CustomColorsCookie, _customColors );
	}

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

	/// <summary>
	/// Wraps color swatches to the available width, with at most 10 per row.
	/// </summary>
	sealed class ColorSwatchPalette : Widget
	{
		const int MaxPerRow = 10;
		const float Spacing = 4f;

		readonly Action<Color> _onPick;
		readonly Action<Color> _onRemove;
		readonly List<(Color Color, string ToolTip, bool Removable)> _entries = new();
		int _builtPerRow = -1;

		public ColorSwatchPalette( Action<Color> onPick, Action<Color> onRemove ) : base( null )
		{
			_onPick = onPick;
			_onRemove = onRemove;
			HorizontalSizeMode = SizeMode.Flexible;
			VerticalSizeMode = SizeMode.CanGrow;
			Layout = Layout.Column();
			Layout.Spacing = Spacing;
		}

		public void SetColors( IEnumerable<(Color Color, string ToolTip, bool Removable)> entries )
		{
			_entries.Clear();
			_entries.AddRange( entries );
			_builtPerRow = -1;
			Rebuild();
		}

		protected override void DoLayout()
		{
			base.DoLayout();

			var perRow = CalcPerRow();
			if ( perRow != _builtPerRow )
				Rebuild();
		}

		int CalcPerRow()
		{
			var cell = Theme.RowHeight + Spacing;
			if ( cell <= 0f || Width <= 1f )
				return MaxPerRow;

			var fit = Math.Max( 1, (int)( (Width + Spacing) / cell ) );
			return Math.Min( MaxPerRow, fit );
		}

		void Rebuild()
		{
			var perRow = CalcPerRow();
			_builtPerRow = perRow;

			Layout.Clear( true );
			Layout.Spacing = Spacing;

			Layout row = null;
			var countInRow = 0;

			foreach ( var entry in _entries )
			{
				if ( row is null || countInRow >= perRow )
				{
					row = Layout.AddRow();
					row.Spacing = Spacing;
					row.Alignment = TextFlag.LeftCenter;
					countInRow = 0;
				}

				var color = entry.Color;
				var removable = entry.Removable;
				var swatch = row.Add( new ColorSwatchButton(
					() => color,
					c => _onPick( c ),
					removable ? () => _onRemove( color ) : null ) );
				swatch.FixedWidth = Theme.RowHeight;
				swatch.FixedHeight = Theme.RowHeight;
				swatch.ToolTip = entry.ToolTip;
				swatch.OpenPickerOnClick = false;
				countInRow++;
			}
		}
	}

	/// <summary>Simple color swatch; left-click opens the engine color picker (optional).</summary>
	sealed class ColorSwatchButton : Widget
	{
		readonly Func<Color> _getColor;
		readonly Action<Color> _setColor;
		readonly Action _onRemove;

		public bool OpenPickerOnClick { get; set; } = true;
		public bool RememberOnClose { get; set; }
		public Action<Color> OnRemember { get; set; }

		public ColorSwatchButton( Func<Color> getColor, Action<Color> setColor, Action onRemove = null ) : base( null )
		{
			_getColor = getColor;
			_setColor = setColor;
			_onRemove = onRemove;
			Cursor = CursorShape.Finger;
			MouseLeftPress = OnClicked;
			MouseRightPress = OnRightClicked;
		}

		void OnClicked()
		{
			if ( !OpenPickerOnClick )
			{
				_setColor( _getColor() );
				return;
			}

			var popup = new ColorPopup( null );
			popup.FixedWidth = 233;
			popup.Layout = Layout.Column();
			popup.Position = Editor.Application.CursorPosition;

			var colorPicker = popup.Layout.Add( new ColorPicker( popup, _getColor() ), 1 );
			colorPicker.OnChildValuesChanged += _ =>
			{
				_setColor( colorPicker.Value );
				Update();
			};

			if ( RememberOnClose )
			{
				popup.Closed = () => OnRemember?.Invoke( colorPicker.Value );
			}

			popup.Show();
			popup.ConstrainToScreen();
		}

		void OnRightClicked()
		{
			_onRemove?.Invoke();
		}

		protected override void OnPaint()
		{
			var color = _getColor();
			const float radius = 3f;
			const float border = 1.25f;

			Paint.Antialiasing = true;

			// Solid white rounded plate first so color can't bleed past the corners.
			Paint.ClearPen();
			Paint.SetBrush( Color.White );
			Paint.DrawRect( LocalRect, radius );

			var inner = LocalRect.Shrink( border );
			var innerRadius = Math.Max( radius - border * 0.5f, 0f );
			Paint.SetBrush( color );
			Paint.DrawRect( inner, innerRadius );
		}
	}

	/// <summary>Popup that notifies when closed so custom colors can be saved once.</summary>
	sealed class ColorPopup : PopupWidget
	{
		public Action Closed { get; set; }

		public ColorPopup( Widget parent ) : base( parent )
		{
		}

		protected override void OnClosed()
		{
			Closed?.Invoke();
			base.OnClosed();
		}
	}
}

/// <summary>Menu + View registration for Sub Hierarchy.</summary>
public static class SubHierarchyMenus
{
	[Menu( "Editor", "Sub Hierarchy/Open", "account_tree" )]
	public static void Open() => SubHierarchyDock.Open();

	[Event( "tools.editorwindow.postcreateview" )]
	static void AddViewMenu( Menu menu )
	{
		menu.AddSeparator();
		menu.AddOption( "Sub Hierarchy", "account_tree", () => SubHierarchyDock.Open() );
	}

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