using System;
using System.Collections.Generic;
using System.Runtime.CompilerServices;
using System.Runtime.InteropServices;
using MoonTools.ECS.Collections;

namespace MoonTools.ECS
{
	internal abstract class RelationStorage
	{
		public abstract unsafe void Set(int entityA, int entityB, void* relationData);
		public abstract int GetStorageIndex(int entityA, int entityB);
		public abstract unsafe void* Get(int relationStorageIndex);
		public abstract void UnrelateAll(int entityID);
		public abstract ReverseSpanEnumerator<(Entity, Entity)> All();
		public abstract RelationStorage CreateStorage();
		public abstract void Clear();
	}

	// Relation is the two entities, A related to B.
	// TRelation is the data attached to the relation.
	internal unsafe class RelationStorage<TRelation> : RelationStorage, IDisposable where TRelation : unmanaged
	{
		private int count = 0;
		private int capacity = 16;
		private (Entity, Entity)* relations;
		private TRelation* relationDatas;
		private Dictionary<(Entity, Entity), int> indices = new Dictionary<(Entity, Entity), int>(16);
		private Dictionary<int, IndexableSet<Entity>> outRelations = new Dictionary<int, IndexableSet<Entity>>(16);
		private Dictionary<int, IndexableSet<Entity>> inRelations = new Dictionary<int, IndexableSet<Entity>>(16);
		private Stack<IndexableSet<Entity>> listPool = new Stack<IndexableSet<Entity>>();

		private bool disposed;

		public RelationStorage()
		{
			relations = ((Entity, Entity)*) NativeMemory.Alloc((nuint) (capacity * Unsafe.SizeOf<(Entity, Entity)>()));
			relationDatas = (TRelation*) NativeMemory.Alloc((nuint) (capacity * Unsafe.SizeOf<TRelation>()));
		}

		public override ReverseSpanEnumerator<(Entity, Entity)> All()
		{
			return new ReverseSpanEnumerator<(Entity, Entity)>(new Span<(Entity, Entity)>(relations, count));
		}

		public void Set(in Entity entityA, in Entity entityB, TRelation relationData)
		{
			var relation = (entityA, entityB);

			if (indices.TryGetValue(relation, out var index))
			{
				relationDatas[index] = relationData;
				return;
			}

			var idA = entityA.ID;
			var idB = entityB.ID;

			if (!outRelations.ContainsKey(idA))
			{
				outRelations[idA] = AcquireHashSetFromPool();
			}
			outRelations[idA].Add(idB);

			if (!inRelations.ContainsKey(idB))
			{
				inRelations[idB] = AcquireHashSetFromPool();
			}
			inRelations[idB].Add(idA);

			if (count >= capacity)
			{
				capacity *= 2;
				relations = ((Entity, Entity)*) NativeMemory.Realloc(relations, (nuint) (capacity * Unsafe.SizeOf<(Entity, Entity)>()));
				relationDatas = (TRelation*) NativeMemory.Realloc(relationDatas, (nuint) (capacity * Unsafe.SizeOf<TRelation>()));
			}

			relations[count] = relation;
			relationDatas[count] = relationData;
			indices.Add(relation, count);
			count += 1;
		}

		public TRelation Get(in Entity entityA, in Entity entityB)
		{
			return relationDatas[indices[(entityA, entityB)]];
		}

		public bool Has((Entity, Entity) relation)
		{
			return indices.ContainsKey(relation);
		}

		public ReverseSpanEnumerator<Entity> OutRelations(int entityID)
		{
			if (outRelations.TryGetValue(entityID, out var entityOutRelations))
			{
				return entityOutRelations.GetEnumerator();
			}
			else
			{
				return ReverseSpanEnumerator<Entity>.Empty;
			}
		}

		public Entity OutFirst(int entityID)
		{
			return OutNth(entityID, 0);
		}

		public Entity OutNth(int entityID, int n)
		{
#if DEBUG
			if (!outRelations.ContainsKey(entityID) || outRelations[entityID].Count == 0)
			{
				throw new KeyNotFoundException("No out relations to this entity!");
			}
#endif
			return outRelations[entityID][n];
		}

		public bool HasOutRelation(int entityID)
		{
			return outRelations.ContainsKey(entityID) && outRelations[entityID].Count > 0;
		}

		public int OutRelationCount(int entityID)
		{
			return outRelations.TryGetValue(entityID, out var entityOutRelations) ? entityOutRelations.Count : 0;
		}

		public ReverseSpanEnumerator<Entity> InRelations(int entityID)
		{
			if (inRelations.TryGetValue(entityID, out var entityInRelations))
			{
				return entityInRelations.GetEnumerator();
			}
			else
			{
				return ReverseSpanEnumerator<Entity>.Empty;
			}
		}

		public Entity InFirst(int entityID)
		{
			return InNth(entityID, 0);
		}

		public Entity InNth(int entityID, int n)
		{
#if DEBUG
			if (!inRelations.ContainsKey(entityID) || inRelations[entityID].Count == 0)
			{
				throw new KeyNotFoundException("No in relations to this entity!");
			}
#endif

			return inRelations[entityID][n];
		}

		public bool HasInRelation(int entityID)
		{
			return inRelations.ContainsKey(entityID) && inRelations[entityID].Count > 0;
		}

		public int InRelationCount(int entityID)
		{
			return inRelations.TryGetValue(entityID, out var entityInRelations) ? entityInRelations.Count : 0;
		}

		public (bool, bool) Remove(in Entity entityA, in Entity entityB)
		{
			var aEmpty = false;
			var bEmpty = false;
			var relation = (entityA, entityB);

			if (outRelations.TryGetValue(entityA.ID, out var entityOutRelations))
			{
				entityOutRelations.Remove(entityB.ID);
				if (outRelations[entityA.ID].Count == 0)
				{
					aEmpty = true;
				}
			}

			if (inRelations.TryGetValue(entityB.ID, out var entityInRelations))
			{
				entityInRelations.Remove(entityA.ID);
				if (inRelations[entityB.ID].Count == 0)
				{
					bEmpty = true;
				}
			}

			if (indices.TryGetValue(relation, out var index))
			{
				var lastElementIndex = count - 1;

				// move an element into the hole
				if (index != lastElementIndex)
				{
					var lastRelation = relations[lastElementIndex];
					indices[lastRelation] = index;
					relationDatas[index] = relationDatas[lastElementIndex];
					relations[index] = lastRelation;
				}

				count -= 1;
				indices.Remove(relation);
			}

			return (aEmpty, bEmpty);
		}

		private IndexableSet<Entity> AcquireHashSetFromPool()
		{
			if (listPool.Count == 0)
			{
				listPool.Push(new IndexableSet<Entity>());
			}

			return listPool.Pop();
		}

		private void ReturnHashSetToPool(IndexableSet<Entity> hashSet)
		{
			hashSet.Clear();
			listPool.Push(hashSet);
		}

		// untyped methods used for internal implementation

		public override unsafe void Set(int entityA, int entityB, void* relationData)
		{
			Set(entityA, entityB, *(TRelation*) relationData);
		}

		public override int GetStorageIndex(int entityA, int entityB)
		{
			return indices[(entityA, entityB)];
		}

		public override unsafe void* Get(int relationStorageIndex)
		{
			return &relationDatas[relationStorageIndex];
		}

		public override void UnrelateAll(int entityID)
		{
			if (outRelations.TryGetValue(entityID, out var entityOutRelations))
			{
				foreach (var entityB in entityOutRelations)
				{
					Remove(entityID, entityB);
				}

				ReturnHashSetToPool(entityOutRelations);
				outRelations.Remove(entityID);
			}

			if (inRelations.TryGetValue(entityID, out var entityInRelations))
			{
				foreach (var entityA in entityInRelations)
				{
					Remove(entityA, entityID);
				}

				ReturnHashSetToPool(entityInRelations);
				inRelations.Remove(entityID);
			}
		}

		public override RelationStorage<TRelation> CreateStorage()
		{
			return new RelationStorage<TRelation>();
		}

		public override void Clear()
		{
			count = 0;
			indices.Clear();

			foreach (var set in inRelations.Values)
			{
				ReturnHashSetToPool(set);
			}
			inRelations.Clear();

			foreach (var set in outRelations.Values)
			{
				ReturnHashSetToPool(set);
			}
			outRelations.Clear();
		}

		protected virtual void Dispose(bool disposing)
		{
			if (!disposed)
			{
				Clear();

				if (disposing)
				{
					foreach (var set in listPool)
					{
						set.Dispose();
					}
				}

				NativeMemory.Free(relations);
				NativeMemory.Free(relationDatas);
				relations = null;
				relationDatas = null;

				disposed = true;
			}
		}

		~RelationStorage()
		{
			// Do not change this code. Put cleanup code in 'Dispose(bool disposing)' method
			Dispose(disposing: false);
		}

		public void Dispose()
		{
			// Do not change this code. Put cleanup code in 'Dispose(bool disposing)' method
			Dispose(disposing: true);
			GC.SuppressFinalize(this);
		}
	}
}