using Editor.Prism.Core;
using Editor.Prism.Model;
using Editor.Prism.Undo;

namespace Editor.Prism.Ui;

/// <summary>Which edge or axis an alignment operation snaps to.</summary>
public enum AlignEdge
{
	/// <summary>Align left edges.</summary>
	Left,
	/// <summary>Align horizontal centres.</summary>
	CenterX,
	/// <summary>Align right edges.</summary>
	Right,
	/// <summary>Align top edges.</summary>
	Top,
	/// <summary>Align vertical middles.</summary>
	Middle,
	/// <summary>Align bottom edges.</summary>
	Bottom
}

/// <summary>Tuning for <see cref="AlignTools.AutoLayout"/>.</summary>
public sealed record AutoLayoutOptions
{
	/// <summary>Horizontal gap between one layer's right edge and the next layer's left edge.</summary>
	public float LayerGap { get; init; } = 96f;

	/// <summary>Vertical gap between two cards in the same layer.</summary>
	public float NodeGap { get; init; } = 32f;

	/// <summary>How many median-ordering sweeps to run. Four is where the returns stop.</summary>
	public int OrderingPasses { get; init; } = 4;

	/// <summary>How many straightening passes to run after ordering.</summary>
	public int StraightenPasses { get; init; } = 6;

	/// <summary>Grid the result snaps to. Zero disables snapping.</summary>
	/// <remarks>
	/// Initialised from <c>PrismConstants.GridSize</c> rather than <c>PrismTheme.GridSize</c> on purpose.
	/// s&amp;box's code generator copies an auto-property's initialiser verbatim into a generated
	/// <c>[DefaultValue( ... )]</c>, so the initialiser of an int/float/bool/enum auto-property has to be a
	/// compile-time constant. <c>PrismTheme.GridSize</c> is a mutable static so a theme can change it,
	/// which would make the generated attribute fail to compile in the editor.
	/// </remarks>
	public float Grid { get; init; } = PrismConstants.GridSize;

	/// <summary>The defaults, which are what the toolbar and the Graph menu use.</summary>
	public static readonly AutoLayoutOptions Default = new();
}

/// <summary>
/// Align, distribute and lay out.
/// <para>
/// Everything here is a pure function of node positions plus a measuring delegate, and every mutation
/// goes through <c>GraphMutations.Move</c> inside one scope — so aligning eleven nodes is one entry in
/// the History panel, and undoing it puts all eleven back.
/// </para>
/// <para>
/// The auto-layout is a real layered layout, not a grid sort: longest-path layering, median-heuristic
/// crossing reduction with adjacent-transpose refinement, then a priority straightening pass that pulls
/// each node towards the average of what it is wired to without letting cards overlap. On a shader
/// graph — wide, shallow, almost always a DAG — that produces something a human would accept.
/// </para>
/// </summary>
public static class AlignTools
{
	/// <summary>Fallback card size when nothing better is known, in scene units.</summary>
	static readonly Vector2 s_defaultSize = new( PrismTheme.NodeMinWidth, 96f );

	// ---------------------------------------------------------------- align ----

	/// <summary>Align a selection to one edge. Returns how many nodes moved.</summary>
	public static int Align( GraphMutations mutations, IEnumerable<PrismNode> nodes, AlignEdge edge,
		Func<PrismNode, Vector2> measure = null )
	{
		var targets = Materialise( nodes );

		if ( mutations is null || targets.Count < 2 ) return 0;

		measure ??= Measure;

		var horizontal = edge is AlignEdge.Left or AlignEdge.CenterX or AlignEdge.Right;
		var anchor = edge switch
		{
			AlignEdge.Left or AlignEdge.Top => 0f,
			AlignEdge.CenterX or AlignEdge.Middle => 0.5f,
			_ => 1f
		};

		var min = float.MaxValue;
		var max = float.MinValue;

		foreach ( var node in targets )
		{
			var size = measure( node );
			var position = node.Position;

			var low = horizontal ? position.x : position.y;
			var high = low + ( horizontal ? size.x : size.y );

			min = MathF.Min( min, low );
			max = MathF.Max( max, high );
		}

		var line = min + ( max - min ) * anchor;
		var moves = new Dictionary<NodeId, Vector2>();

		foreach ( var node in targets )
		{
			var size = measure( node );
			var position = node.Position;

			if ( horizontal ) position.x = line - size.x * anchor;
			else position.y = line - size.y * anchor;

			moves[node.Id] = Snap( position, PrismTheme.GridSize );
		}

		return Apply( mutations, moves, "Align" );
	}

	// ---------------------------------------------------------------- distribute ----

	/// <summary>
	/// Space a selection evenly between its two extremes, keeping the outermost nodes where they are.
	/// Gaps are equalised, not centres, so cards of different heights read as evenly spaced.
	/// </summary>
	public static int Distribute( GraphMutations mutations, IEnumerable<PrismNode> nodes, bool horizontal,
		Func<PrismNode, Vector2> measure = null )
	{
		var targets = Materialise( nodes );

		if ( mutations is null || targets.Count < 3 ) return 0;

		measure ??= Measure;

		var ordered = targets
			.OrderBy( x => horizontal ? x.Position.x : x.Position.y )
			.ToList();

		var first = ordered[0];
		var last = ordered[^1];

		var start = horizontal ? first.Position.x : first.Position.y;
		var end = horizontal
			? last.Position.x + measure( last ).x
			: last.Position.y + measure( last ).y;

		var occupied = 0f;

		foreach ( var node in ordered )
		{
			var size = measure( node );

			occupied += horizontal ? size.x : size.y;
		}

		var gap = ( end - start - occupied ) / ( ordered.Count - 1 );
		var moves = new Dictionary<NodeId, Vector2>();
		var cursor = start;

		foreach ( var node in ordered )
		{
			var size = measure( node );
			var position = node.Position;

			if ( horizontal ) position.x = cursor;
			else position.y = cursor;

			moves[node.Id] = Snap( position, PrismTheme.GridSize );

			cursor += ( horizontal ? size.x : size.y ) + gap;
		}

		return Apply( mutations, moves, "Distribute" );
	}

	/// <summary>
	/// Snap a selection — or the whole document — to the grid. A non-positive grid means "the theme's",
	/// which cannot be a default parameter value because the grid size is a mutable theme token.
	/// </summary>
	public static int SnapToGrid( GraphMutations mutations, IEnumerable<PrismNode> nodes,
		float grid = -1f )
	{
		if ( grid <= 0f ) grid = PrismTheme.GridSize;

		var targets = Materialise( nodes );

		if ( mutations is null || targets.Count == 0 || grid <= 0f ) return 0;

		var moves = new Dictionary<NodeId, Vector2>();

		foreach ( var node in targets )
		{
			moves[node.Id] = Snap( node.Position, grid );
		}

		return Apply( mutations, moves, "Snap To Grid" );
	}

	// ---------------------------------------------------------------- auto-layout ----

	/// <summary>
	/// Lay out a graph — or a subset of it — left to right by data flow.
	/// <para>
	/// Nodes outside <paramref name="subset"/> are not moved and do not participate, so auto-laying a
	/// selection rearranges only that selection, in place, within its own bounding box.
	/// </para>
	/// </summary>
	public static int AutoLayout( GraphMutations mutations, PrismGraph graph, IEnumerable<PrismNode> subset = null,
		AutoLayoutOptions options = null, Func<PrismNode, Vector2> measure = null )
	{
		if ( mutations is null || graph is null ) return 0;

		options ??= AutoLayoutOptions.Default;
		measure ??= Measure;

		var nodes = Materialise( subset );

		if ( nodes.Count == 0 ) nodes = graph.Nodes.Where( x => x is not null ).ToList();
		if ( nodes.Count < 2 ) return 0;

		var index = new Dictionary<NodeId, int>();

		for ( int i = 0; i < nodes.Count; i++ ) index[nodes[i].Id] = i;

		var successors = BuildAdjacency( graph, nodes, index, out var predecessors );

		// Where the block starts now, so the result lands roughly where the user was looking rather
		// than at the origin.
		var origin = Origin( nodes );

		var layers = Layer( nodes, predecessors );
		var columns = GroupIntoColumns( nodes, layers );

		OrderColumns( columns, successors, predecessors, options.OrderingPasses );

		var moves = Place( columns, nodes, successors, predecessors, measure, options, origin );

		return Apply( mutations, moves, "Auto-Layout" );
	}

	/// <summary>Nodes that feed nothing reachable from an output. The "Clean Up Unused" candidates.</summary>
	public static IReadOnlyList<PrismNode> UnusedNodes( PrismGraph graph )
	{
		if ( graph is null ) return Array.Empty<PrismNode>();

		var reachable = PrismLog.Guard<IReadOnlyCollection<NodeId>>( "Find reachable nodes",
			() => GraphQueries.ReachableFromOutputs( graph ), Array.Empty<NodeId>() );

		if ( reachable is null || reachable.Count == 0 ) return Array.Empty<PrismNode>();

		var keep = new HashSet<NodeId>( reachable );

		return graph.Nodes
			.Where( x => x is not null && !keep.Contains( x.Id ) )
			.ToArray();
	}

	// ---------------------------------------------------------------- layering ----

	static Dictionary<int, List<int>> BuildAdjacency( PrismGraph graph, IReadOnlyList<PrismNode> nodes,
		IReadOnlyDictionary<NodeId, int> index, out Dictionary<int, List<int>> predecessors )
	{
		var successors = new Dictionary<int, List<int>>();

		predecessors = new Dictionary<int, List<int>>();

		for ( int i = 0; i < nodes.Count; i++ )
		{
			successors[i] = new List<int>();
			predecessors[i] = new List<int>();
		}

		foreach ( var edge in graph.Edges )
		{
			if ( edge is null ) continue;
			if ( !index.TryGetValue( edge.FromNode, out var from ) ) continue;
			if ( !index.TryGetValue( edge.ToNode, out var to ) ) continue;
			if ( from == to ) continue;

			if ( !successors[from].Contains( to ) ) successors[from].Add( to );
			if ( !predecessors[to].Contains( from ) ) predecessors[to].Add( from );
		}

		return successors;
	}

	/// <summary>
	/// Longest-path layering. Cycles cannot make this diverge: the depth-first walk carries a visiting
	/// set and treats a back edge as absent, which is exactly how a cyclic graph should lay out — as
	/// though the loop had been cut once.
	/// </summary>
	static int[] Layer( IReadOnlyList<PrismNode> nodes, IReadOnlyDictionary<int, List<int>> predecessors )
	{
		var layers = new int[nodes.Count];
		var state = new byte[nodes.Count];

		for ( int i = 0; i < nodes.Count; i++ ) layers[i] = -1;

		for ( int i = 0; i < nodes.Count; i++ )
		{
			Depth( i );
		}

		return layers;

		int Depth( int node )
		{
			if ( layers[node] >= 0 ) return layers[node];

			// Already on the stack: a cycle. Report zero so the caller carries on without recursing.
			if ( state[node] == 1 ) return 0;

			state[node] = 1;

			var depth = 0;

			foreach ( var predecessor in predecessors[node] )
			{
				depth = Math.Max( depth, Depth( predecessor ) + 1 );
			}

			state[node] = 2;
			layers[node] = depth;

			return depth;
		}
	}

	static List<List<int>> GroupIntoColumns( IReadOnlyList<PrismNode> nodes, int[] layers )
	{
		var count = 0;

		for ( int i = 0; i < layers.Length; i++ ) count = Math.Max( count, layers[i] + 1 );

		var columns = new List<List<int>>( Math.Max( 1, count ) );

		for ( int i = 0; i < count; i++ ) columns.Add( new List<int>() );

		// Seeding each column in current vertical order means a graph that was already tidy stays close
		// to where it was, which makes auto-layout feel like a nudge rather than a shuffle.
		var order = Enumerable.Range( 0, nodes.Count )
			.OrderBy( i => nodes[i].Position.y )
			.ThenBy( i => nodes[i].Position.x );

		foreach ( var i in order )
		{
			var layer = Math.Clamp( layers[i], 0, count - 1 );

			columns[layer].Add( i );
		}

		return columns;
	}

	// ---------------------------------------------------------------- ordering ----

	static void OrderColumns( List<List<int>> columns, IReadOnlyDictionary<int, List<int>> successors,
		IReadOnlyDictionary<int, List<int>> predecessors, int passes )
	{
		var rank = new Dictionary<int, int>();

		Rerank();

		for ( int pass = 0; pass < Math.Max( 1, passes ); pass++ )
		{
			var forward = ( pass & 1 ) == 0;

			if ( forward )
			{
				for ( int c = 1; c < columns.Count; c++ ) Median( columns[c], predecessors );
			}
			else
			{
				for ( int c = columns.Count - 2; c >= 0; c-- ) Median( columns[c], successors );
			}

			Rerank();

			foreach ( var column in columns ) Transpose( column );

			Rerank();
		}

		void Rerank()
		{
			rank.Clear();

			foreach ( var column in columns )
			{
				for ( int i = 0; i < column.Count; i++ ) rank[column[i]] = i;
			}
		}

		void Median( List<int> column, IReadOnlyDictionary<int, List<int>> neighbours )
		{
			if ( column.Count < 2 ) return;

			var keys = new Dictionary<int, float>( column.Count );

			for ( int i = 0; i < column.Count; i++ )
			{
				keys[column[i]] = MedianOf( column[i], neighbours, rank, i );
			}

			column.Sort( ( a, b ) => keys[a].CompareTo( keys[b] ) );
		}

		void Transpose( List<int> column )
		{
			if ( column.Count < 2 ) return;

			var improved = true;
			var guard = 0;

			while ( improved && guard++ < 8 )
			{
				improved = false;

				for ( int i = 0; i + 1 < column.Count; i++ )
				{
					var a = column[i];
					var b = column[i + 1];

					if ( Crossings( a, b ) <= Crossings( b, a ) ) continue;

					column[i] = b;
					column[i + 1] = a;
					rank[b] = i;
					rank[a] = i + 1;
					improved = true;
				}
			}
		}

		int Crossings( int upper, int lower )
		{
			var count = 0;

			foreach ( var a in Neighbours( upper ) )
			{
				foreach ( var b in Neighbours( lower ) )
				{
					if ( Position( a ) > Position( b ) ) count++;
				}
			}

			return count;
		}

		IEnumerable<int> Neighbours( int node ) =>
			predecessors[node].Concat( successors[node] );

		int Position( int node ) => rank.TryGetValue( node, out var value ) ? value : 0;
	}

	static float MedianOf( int node, IReadOnlyDictionary<int, List<int>> neighbours,
		IReadOnlyDictionary<int, int> rank, int fallback )
	{
		if ( !neighbours.TryGetValue( node, out var list ) || list.Count == 0 ) return fallback;

		var positions = new List<int>( list.Count );

		foreach ( var neighbour in list )
		{
			if ( rank.TryGetValue( neighbour, out var value ) ) positions.Add( value );
		}

		if ( positions.Count == 0 ) return fallback;

		positions.Sort();

		var middle = positions.Count / 2;

		return positions.Count % 2 == 1
			? positions[middle]
			: ( positions[middle - 1] + positions[middle] ) * 0.5f;
	}

	// ---------------------------------------------------------------- placement ----

	static Dictionary<NodeId, Vector2> Place( List<List<int>> columns, IReadOnlyList<PrismNode> nodes,
		IReadOnlyDictionary<int, List<int>> successors, IReadOnlyDictionary<int, List<int>> predecessors,
		Func<PrismNode, Vector2> measure, AutoLayoutOptions options, Vector2 origin )
	{
		var sizes = new Vector2[nodes.Count];
		var tops = new float[nodes.Count];

		// Column x positions, packed left to right by the widest card in each column.
		var x = origin.x;
		var columnX = new float[columns.Count];

		for ( int c = 0; c < columns.Count; c++ )
		{
			columnX[c] = x;

			var widest = PrismTheme.NodeMinWidth;

			foreach ( var slot in columns[c] )
			{
				var size = SizeOf( slot );

				widest = MathF.Max( widest, size.x );
			}

			x += widest + options.LayerGap;
		}

		// First pass: stack each column from the top.
		foreach ( var column in columns )
		{
			var y = origin.y;

			foreach ( var slot in column )
			{
				tops[slot] = y;
				y += SizeOf( slot ).y + options.NodeGap;
			}
		}

		// Straightening: pull every node towards the average centre of what it is wired to, then push
		// the column apart again so nothing overlaps. Alternating direction keeps long chains straight
		// instead of dragging the whole graph one way.
		for ( int pass = 0; pass < Math.Max( 0, options.StraightenPasses ); pass++ )
		{
			var forward = ( pass & 1 ) == 0;

			for ( int i = 0; i < columns.Count; i++ )
			{
				var c = forward ? i : columns.Count - 1 - i;
				var column = columns[c];

				if ( column.Count == 0 ) continue;

				foreach ( var slot in column )
				{
					var wanted = Barycentre( slot, forward );

					if ( float.IsNaN( wanted ) ) continue;

					tops[slot] = wanted - SizeOf( slot ).y * 0.5f;
				}

				Separate( column );
			}
		}

		var moves = new Dictionary<NodeId, Vector2>();

		for ( int c = 0; c < columns.Count; c++ )
		{
			foreach ( var slot in columns[c] )
			{
				var position = new Vector2( columnX[c], tops[slot] );

				moves[nodes[slot].Id] = Snap( position, options.Grid );
			}
		}

		return moves;

		Vector2 SizeOf( int slot )
		{
			if ( sizes[slot] != default ) return sizes[slot];

			var node = nodes[slot];
			var size = node is null ? s_defaultSize : measure( node );

			if ( size.x <= 1f || size.y <= 1f ) size = s_defaultSize;

			sizes[slot] = size;

			return size;
		}

		float Barycentre( int slot, bool fromPredecessors )
		{
			var neighbours = fromPredecessors ? predecessors[slot] : successors[slot];

			if ( neighbours.Count == 0 ) return float.NaN;

			var total = 0f;

			foreach ( var neighbour in neighbours )
			{
				total += tops[neighbour] + SizeOf( neighbour ).y * 0.5f;
			}

			return total / neighbours.Count;
		}

		void Separate( List<int> column )
		{
			// Order is fixed by the crossing-reduction pass, so separation only has to push downwards
			// from the first node and then lift the whole column back to where its centre wanted to be.
			var before = 0f;
			var after = 0f;

			foreach ( var slot in column ) before += tops[slot];

			for ( int i = 1; i < column.Count; i++ )
			{
				var previous = column[i - 1];
				var current = column[i];
				var floor = tops[previous] + SizeOf( previous ).y + options.NodeGap;

				if ( tops[current] < floor ) tops[current] = floor;
			}

			foreach ( var slot in column ) after += tops[slot];

			var drift = ( after - before ) / column.Count;

			if ( MathF.Abs( drift ) < 0.01f ) return;

			foreach ( var slot in column ) tops[slot] -= drift;
		}
	}

	// ---------------------------------------------------------------- helpers ----

	static List<PrismNode> Materialise( IEnumerable<PrismNode> nodes ) =>
		nodes is null
			? new List<PrismNode>()
			: nodes.Where( x => x is not null && x.Id.IsValid ).Distinct().ToList();

	static Vector2 Origin( IReadOnlyList<PrismNode> nodes )
	{
		var min = new Vector2( float.MaxValue, float.MaxValue );

		foreach ( var node in nodes )
		{
			min.x = MathF.Min( min.x, node.Position.x );
			min.y = MathF.Min( min.y, node.Position.y );
		}

		return min.x == float.MaxValue ? Vector2.Zero : min;
	}

	static Vector2 Snap( Vector2 position, float grid )
	{
		if ( grid <= 0f ) return position;

		return new Vector2(
			MathF.Round( position.x / grid ) * grid,
			MathF.Round( position.y / grid ) * grid );
	}

	/// <summary>
	/// A card size when the canvas is not around to ask. Derived from the port count, which is what
	/// actually drives card height, so a headless layout is still roughly right.
	/// </summary>
	public static Vector2 Measure( PrismNode node )
	{
		if ( node is null ) return s_defaultSize;

		var rows = Math.Max( node.Inputs?.Count ?? 0, node.Outputs?.Count ?? 0 );
		var height = PrismTheme.NodeHeaderHeight + Math.Max( 1, rows ) * PrismTheme.PortPitch + PrismTheme.Rhythm;

		if ( ( node.Flags & NodeFlags.Preview ) != 0 ) height += PrismTheme.ThumbnailSize;
		if ( ( node.Flags & NodeFlags.Collapsed ) != 0 ) height = PrismTheme.NodeHeaderHeight;

		var title = node.Descriptor?.Title ?? string.Empty;
		var width = Math.Clamp( PrismTheme.NodeMinWidth + title.Length * 2f,
			PrismTheme.NodeMinWidth, PrismTheme.NodeMaxWidth );

		return new Vector2( width, height );
	}

	static int Apply( GraphMutations mutations, IReadOnlyDictionary<NodeId, Vector2> moves, string label )
	{
		if ( mutations is null || moves is null || moves.Count == 0 ) return 0;

		using ( mutations.Begin( label ) )
		{
			return mutations.Move( moves, label );
		}
	}
}