using System.Collections.Generic;
using MoonWorks.Math.Fixed;

namespace MoonWorks.Collision.Fixed
{
	/// <summary>
	/// Used to quickly check if two shapes are potentially overlapping.
	/// </summary>
	/// <typeparam name="T">The type that will be used to uniquely identify shape-transform pairs.</typeparam>
	public class SpatialHash2D<T, U> where T : struct, System.IEquatable<T> where U : IHasAABB2D
	{
		private readonly Fix64 cellSize;

		private readonly Dictionary<long, HashSet<T>> hashDictionary = new Dictionary<long, HashSet<T>>();
		private readonly Dictionary<T, (U, Transform2D, uint)> IDLookup = new Dictionary<T, (U, Transform2D, uint)>();

		public int MinX { get; private set; } = 0;
		public int MaxX { get; private set; } = 0;
		public int MinY { get; private set; } = 0;
		public int MaxY { get; private set; } = 0;

		private Queue<HashSet<T>> hashSetPool = new Queue<HashSet<T>>();

		public SpatialHash2D(int cellSize)
		{
			this.cellSize = new Fix64(cellSize);
		}

		private (int, int) Hash(Vector2 position)
		{
			return ((int) (position.X / cellSize), (int) (position.Y / cellSize));
		}

		/// <summary>
		/// Inserts an element into the SpatialHash.
		/// </summary>
		/// <param name="id">A unique ID for the shape-transform pair.</param>
		/// <param name="shape"></param>
		/// <param name="transform2D"></param>
		/// <param name="collisionGroups">A bitmask value specifying the groups this object belongs to.</param>
		public void Insert(T id, U shape, Transform2D transform2D, uint collisionGroups = uint.MaxValue)
		{
			var box = AABB2D.Transformed(shape.AABB, transform2D);
			var minHash = Hash(box.Min);
			var maxHash = Hash(box.Max);

			foreach (var key in Keys(minHash.Item1, minHash.Item2, maxHash.Item1, maxHash.Item2))
			{
				if (!hashDictionary.ContainsKey(key))
				{
					hashDictionary.Add(key, new HashSet<T>());
				}

				hashDictionary[key].Add(id);
				IDLookup[id] = (shape, transform2D, collisionGroups);
			}

			MinX = System.Math.Min(MinX, minHash.Item1);
			MinY = System.Math.Min(MinY, minHash.Item2);
			MaxX = System.Math.Max(MaxX, maxHash.Item1);
			MaxY = System.Math.Max(MaxY, maxHash.Item2);
		}

		/// <summary>
		/// Retrieves all the potential collisions of a shape-transform pair. Excludes any shape-transforms with the given ID.
		/// </summary>
		public RetrieveEnumerator Retrieve<V>(T id, V hasAABB, Transform2D transform2D, uint collisionMask = uint.MaxValue) where V : IHasAABB2D
		{
			var box = AABB2D.Transformed(hasAABB.AABB, transform2D);
			var (minX, minY) = Hash(box.Min);
			var (maxX, maxY) = Hash(box.Max);

			if (minX < MinX) { minX = MinX; }
			if (maxX > MaxX) { maxX = MaxX; }
			if (minY < MinY) { minY = MinY; }
			if (maxY > MaxY) { maxY = MaxY; }

			return new RetrieveEnumerator(
				this,
				Keys(minX, minY, maxX, maxY),
				id,
				collisionMask
			);
		}

		/// <summary>
		/// Retrieves all the potential collisions of a shape-transform pair.
		/// </summary>
		public RetrieveEnumerator Retrieve<V>(V hasAABB, Transform2D transform2D, uint collisionMask = uint.MaxValue) where V : IHasAABB2D
		{
			var box = AABB2D.Transformed(hasAABB.AABB, transform2D);
			return Retrieve(box, collisionMask);
		}

		/// <summary>
		/// Retrieves objects based on a pre-transformed AABB.
		/// </summary>
		/// <param name="aabb">A transformed AABB.</param>
		/// <returns></returns>
		public RetrieveEnumerator Retrieve(AABB2D aabb, uint collisionMask = uint.MaxValue)
		{
			var (minX, minY) = Hash(aabb.Min);
			var (maxX, maxY) = Hash(aabb.Max);

			if (minX < MinX) { minX = MinX; }
			if (maxX > MaxX) { maxX = MaxX; }
			if (minY < MinY) { minY = MinY; }
			if (maxY > MaxY) { maxY = MaxY; }

			return new RetrieveEnumerator(
				this,
				Keys(minX, minY, maxX, maxY),
				collisionMask
			);
		}

		public void Update(T id, U shape, Transform2D transform2D, uint collisionGroups = uint.MaxValue)
		{
			Remove(id);
			Insert(id, shape, transform2D, collisionGroups);
		}

		/// <summary>
		/// Removes a specific ID from the SpatialHash.
		/// </summary>
		public void Remove(T id)
		{
			var (shape, transform, collisionGroups) = IDLookup[id];

			var box = AABB2D.Transformed(shape.AABB, transform);
			var minHash = Hash(box.Min);
			var maxHash = Hash(box.Max);

			foreach (var key in Keys(minHash.Item1, minHash.Item2, maxHash.Item1, maxHash.Item2))
			{
				if (hashDictionary.ContainsKey(key))
				{
					hashDictionary[key].Remove(id);
				}
			}

			IDLookup.Remove(id);
		}

		/// <summary>
		/// Removes everything that has been inserted into the SpatialHash.
		/// </summary>
		public void Clear()
		{
			foreach (var hash in hashDictionary.Values)
			{
				hash.Clear();
			}

			IDLookup.Clear();
		}

		private static long MakeLong(int left, int right)
		{
			return ((long) left << 32) | ((uint) right);
		}

		internal static KeysEnumerator Keys(int minX, int minY, int maxX, int maxY)
		{
			return new KeysEnumerator(minX, minY, maxX, maxY);
		}

		internal HashSet<T> AcquireHashSet()
		{
			if (hashSetPool.Count == 0)
			{
				hashSetPool.Enqueue(new HashSet<T>());
			}

			var hashSet = hashSetPool.Dequeue();
			hashSet.Clear();
			return hashSet;
		}

		internal void FreeHashSet(HashSet<T> hashSet)
		{
			hashSetPool.Enqueue(hashSet);
		}

		internal ref struct KeysEnumerator
		{
			private int MinX;
			private int MinY;
			private int MaxX;
			private int MaxY;
			private int i, j;

			public KeysEnumerator GetEnumerator() => this;

			public KeysEnumerator(int minX, int minY, int maxX, int maxY)
			{
				MinX = minX;
				MinY = minY;
				MaxX = maxX;
				MaxY = maxY;
				i = minX;
				j = minY - 1;
			}

			public bool MoveNext()
			{
				if (j < MaxY)
				{
					j += 1;
					return true;
				}
				else if (i < MaxX)
				{
					i += 1;
					j = MinY;
					return true;
				}

				return false;
			}

			public long Current
			{
				get
				{
					return MakeLong(i, j);
				}
			}
		}

		public ref struct RetrieveEnumerator
		{
			public SpatialHash2D<T, U> SpatialHash;
			private KeysEnumerator KeysEnumerator;
			private HashSet<T>.Enumerator HashSetEnumerator;
			private bool HashSetEnumeratorActive;
			private HashSet<T> Duplicates;
			private T? ID;
			private uint CollisionMask;

			public RetrieveEnumerator GetEnumerator() => this;

			internal RetrieveEnumerator(
				SpatialHash2D<T, U> spatialHash,
				KeysEnumerator keysEnumerator,
				T id,
				uint collisionMask
			) {
				SpatialHash = spatialHash;
				KeysEnumerator = keysEnumerator;
				HashSetEnumerator = default;
				HashSetEnumeratorActive = false;
				Duplicates = SpatialHash.AcquireHashSet();
				ID = id;
				CollisionMask = collisionMask;
			}

			internal RetrieveEnumerator(
				SpatialHash2D<T, U> spatialHash,
				KeysEnumerator keysEnumerator,
				uint collisionMask
			) {
				SpatialHash = spatialHash;
				KeysEnumerator = keysEnumerator;
				HashSetEnumerator = default;
				HashSetEnumeratorActive = false;
				Duplicates = SpatialHash.AcquireHashSet();
				ID = null;
				CollisionMask = collisionMask;
			}

			public bool MoveNext()
			{
				if (!HashSetEnumeratorActive || !HashSetEnumerator.MoveNext())
				{
					if (!KeysEnumerator.MoveNext())
					{
						SpatialHash.FreeHashSet(Duplicates);
						return false;
					}

					if (SpatialHash.hashDictionary.TryGetValue(KeysEnumerator.Current, out var hashset))
					{
						HashSetEnumerator = hashset.GetEnumerator();
						HashSetEnumeratorActive = true;
					}

					return MoveNext();
				}

				// conditions
				var t = HashSetEnumerator.Current;
				var collisionGroups = SpatialHash.IDLookup[t].Item3;

				if (Duplicates.Contains(t))
				{
					return MoveNext();
				}

				if (ID.HasValue)
				{
					if (ID.Value.Equals(t) || (CollisionMask & collisionGroups) == 0)
					{
						return MoveNext();
					}
				}

				Duplicates.Add(t);
				return true;
			}

			public (T, U, Transform2D, uint) Current
			{
				get
				{
					var t = HashSetEnumerator.Current;
					var (u, transform, groups) = SpatialHash.IDLookup[t];
					return (t, u, transform, groups);
				}
			}
		}
	}
}