using System;
using System.Collections.Generic;
using System.Linq;

namespace Sandbox
{
	/// <summary>
	/// Generate an ordering based on a set of first-most and last-most items, and
	/// individual constraints between pairs of items. All first-most items will be
	/// ordered before all last-most items, and any other items will be put in the
	/// middle unless forced to be elsewhere by a constraint.
	/// </summary>
	internal class SortingHelper
	{
		public struct SortConstraint : IEquatable<SortConstraint>
		{
			public readonly int EarlierIndex;
			public readonly int LaterIndex;

			private readonly int _hashCode;

			public bool IsZero => EarlierIndex == 0 && LaterIndex == 0;

			public SortConstraint Complement => new SortConstraint( LaterIndex, EarlierIndex );

			public SortConstraint( int earlierIndex, int laterIndex )
			{
				EarlierIndex = earlierIndex;
				LaterIndex = laterIndex;

				unchecked
				{
					_hashCode = (EarlierIndex * 397) ^ LaterIndex;
				}
			}

			public bool Equals( SortConstraint other ) => EarlierIndex == other.EarlierIndex && LaterIndex == other.LaterIndex;

			public override bool Equals( object obj ) => obj is SortConstraint other && Equals( other );

			public override int GetHashCode() => _hashCode;
		}

		private readonly int _itemCount;

		private readonly HashSet<SortConstraint> _initialConstraints = new HashSet<SortConstraint>();

		private readonly HashSet<int> _first = new HashSet<int>();
		private readonly HashSet<int> _last = new HashSet<int>();

		public SortingHelper( int itemCount )
		{
			_itemCount = itemCount;
		}

		public void AddConstraint( int earlierIndex, int laterIndex )
		{
			_initialConstraints.Add( new SortConstraint( earlierIndex, laterIndex ) );
		}

		public void AddFirst( int earlierIndex )
		{
			_first.Add( earlierIndex );
		}

		public void AddLast( int laterIndex )
		{
			_last.Add( laterIndex );
		}

		public bool Sort( List<int> result, out SortConstraint invalidConstraint )
		{
			var middle = new HashSet<int>();

			for ( var index = 0; index < _itemCount; ++index )
			{
				if ( !_first.Contains( index ) && !_last.Contains( index ) )
					middle.Add( index );
			}

			var allConstraints = new HashSet<SortConstraint>();
			var newConstraints = new Queue<SortConstraint>();
			var beforeDict = new Dictionary<int, HashSet<int>>();
			var afterDict = new Dictionary<int, HashSet<int>>();

			bool AddWorkingConstraint( int earlierIndex, int laterIndex, out SortConstraint constraint )
			{
				constraint = new SortConstraint( earlierIndex, laterIndex );

				if ( allConstraints.Contains( constraint.Complement ) )
					return false;

				if ( !allConstraints.Add( constraint ) )
					return true;

				newConstraints.Enqueue( constraint );

				if ( !beforeDict.TryGetValue( earlierIndex, out var before ) )
					beforeDict.Add( earlierIndex, before = new HashSet<int>() );

				if ( !afterDict.TryGetValue( laterIndex, out var after ) )
					afterDict.Add( laterIndex, after = new HashSet<int>() );

				before.Add( laterIndex );
				after.Add( earlierIndex );

				return true;
			}

			// Add initial constraints

			foreach ( var initialConstraint in _initialConstraints )
			{
				if ( !AddWorkingConstraint( initialConstraint.EarlierIndex, initialConstraint.LaterIndex, out invalidConstraint ) )
					return false;
			}

			// Everything in _first should be before everything in _last

			foreach ( var earlierIndex in _first )
			{
				foreach ( var laterIndex in _last )
				{
					if ( !AddWorkingConstraint( earlierIndex, laterIndex, out invalidConstraint ) )
						return false;
				}
			}

			// Keep propagating constraints until nothing changes

			while ( newConstraints.TryDequeue( out var nextConstraint ) )
			{
				// if a < b, and b < c, then a < c etc

				if ( beforeDict.TryGetValue( nextConstraint.LaterIndex, out var before ) )
				{
					foreach ( var laterIndex in before )
					{
						if ( !AddWorkingConstraint( nextConstraint.EarlierIndex, laterIndex, out invalidConstraint ) )
							return false;
					}
				}

				if ( afterDict.TryGetValue( nextConstraint.EarlierIndex, out var after ) )
				{
					foreach ( var earlierIndex in after )
					{
						if ( !AddWorkingConstraint( earlierIndex, nextConstraint.LaterIndex, out invalidConstraint ) )
						{
							return false;
						}
					}
				}
			}

			// Now if we have any items that aren't using GroupOrder.First, and haven't
			// determined that they are ordered before another item with GroupOrder.First,
			// we can safely order them after all GroupOrder.First items. And vice versa.

			foreach ( var middleIndex in middle )
			{
				var isBeforeAnyFirst = beforeDict.TryGetValue( middleIndex, out var before )
					&& before.Any( x => _first.Contains( x ) );

				var isAfterAnyLast = afterDict.TryGetValue( middleIndex, out var after )
					&& after.Any( x => _last.Contains( x ) );

				if ( !isBeforeAnyFirst )
				{
					foreach ( var earlierIndex in _first )
						AddWorkingConstraint( earlierIndex, middleIndex, out invalidConstraint );
				}

				if ( !isAfterAnyLast )
				{
					foreach ( var laterIndex in _last )
						AddWorkingConstraint( middleIndex, laterIndex, out invalidConstraint );
				}
			}

			// Now lets add items to the final ordering if all items that should be sorted
			// before them are already added to that ordering. We'll implement this by choosing
			// items that have an empty list / don't appear in afterDict, and update that
			// dictionary as we go.

			var earliestRemaining = new Queue<int>();

			// First, seed the queue with everything that's already not ordered after anything

			for ( var index = 0; index < _itemCount; ++index )
			{
				if ( !afterDict.ContainsKey( index ) )
				{
					earliestRemaining.Enqueue( index );
				}
			}

			result.Clear();

			while ( earliestRemaining.TryDequeue( out var nextIndex ) )
			{
				result.Add( nextIndex );

				foreach ( var laterIndex in beforeDict.TryGetValue( nextIndex, out var laterIndices )
					? laterIndices : Enumerable.Empty<int>() )
				{
					var beforeLater = afterDict[laterIndex];
					beforeLater.Remove( nextIndex );

					if ( beforeLater.Count == 0 )
						earliestRemaining.Enqueue( laterIndex );
				}
			}

			invalidConstraint = default;
			return result.Count == _itemCount;
		}
	}
}
