EntityHashGrid

using SBP;
using SBP.LINK;
using SBP.Meshing;
using System;
using System.Collections;
using System.Collections.Generic;
using Unity.Collections;
using Unity.Collections.LowLevel.Unsafe;
using Unity.Jobs;
using Unity.Mathematics;
using UnityEngine;


// wraps a pointer to a hashset so we can nest in a native container without unity complaining
// could use without the pointer but would need to be careful that the internal state of the struct is stored, you can't pass it around
public unsafe struct UnsafeHashSetHolder<T> : IDisposable , IEnumerable<T>
        where T : unmanaged, IEquatable<T>
{
    // TODO disable ptr restriction
    public UnsafeHashSet<T>* data;

    /// <summary>
    /// Removes all values.
    /// </summary>
    /// <remarks>Does not change the capacity.</remarks>
    public void Clear() => data->Clear();

    /// <summary>
    /// Adds a new value (unless it is already present).
    /// </summary>
    /// <param name="item">The value to add.</param>
    /// <returns>True if the value was not already present.</returns>
    public bool Add(T item) => data->Add(item);

    /// <summary>
    /// Removes a particular value.
    /// </summary>
    /// <param name="item">The value to remove.</param>
    /// <returns>True if the value was present.</returns>
    public bool Remove(T item) => data->Remove(item);

    /// <summary>
    /// Returns true if a particular value is present.
    /// </summary>
    /// <param name="item">The value to check for.</param>
    /// <returns>True if the value was present.</returns>
    public bool Contains(T item) => data->Contains(item);

    /// <summary>
    /// Sets the capacity to match what it would be if it had been originally initialized with all its entries.
    /// </summary>
    public void TrimExcess() => data->TrimExcess();

    /// <summary>
    /// Returns an enumerator over the values of this set.
    /// </summary>
    /// <returns>An enumerator over the values of this set.</returns>
    public UnsafeHashSet<T>.Enumerator GetEnumerator()
    {
        return data->GetEnumerator();
    }

    /// <summary>
    /// This method is not implemented. Use <see cref="GetEnumerator"/> instead.
    /// </summary>
    /// <returns>Throws NotImplementedException.</returns>
    /// <exception cref="NotImplementedException">Method is not implemented.</exception>
    IEnumerator<T> IEnumerable<T>.GetEnumerator()
    {
        throw new NotImplementedException();
    }

    /// <summary>
    /// This method is not implemented. Use <see cref="GetEnumerator"/> instead.
    /// </summary>
    /// <returns>Throws NotImplementedException.</returns>
    /// <exception cref="NotImplementedException">Method is not implemented.</exception>
    IEnumerator IEnumerable.GetEnumerator()
    {
        throw new NotImplementedException();
    }
}

public static class JobsExtensions
{
    public static Allocator Temp => JobsUtility.IsExecutingJob() ? Allocator.Temp : Allocator.TempJob;

    public static void RunSafe<T>(this T job) where T : IJob
    {
        if (JobsUtility.IsExecutingJob())
            job.Execute();
        else
            job.Run();
    }


    public static Allocator AsJobSafe(this Allocator allocator)
    {
        JobsUtility.IsExecutingJob() ? Allocator.Temp : allocator == Allocator.Temp : Allocator.TempJob : allocator;
    }
}
public struct EntityHashGrid : IDisposable
{
    // might be more performant to use int2 xz grid, the majority of air and ground chunks will be empty
    // this will reduce iterating the grid from n^3 to n^2, saving performance on large queries and reducing the
    // performance cost of smaller chunkSizes
    public NativeHashMap<int2, UnsafeHashSetHolder<uint>> entitiesPerChunk;
    public NativeParallelMultiHashMap<uint, int2> chunksPerEntity;
    public NativeHashMap<uint, AABB> boundsPerEntity;

    public uint ecsId;
    // not clear how large this should be. smaller values will result in more accurate results at the expense of
    // more lookups and more duplication.
    public float chunkSize;

    public EntityHashGrid(uint ecsId, float chunkSize)
    {
        entitiesPerChunk = new(32, Allocator.Persistent);
        chunksPerEntity = new(32, Allocator.Persistent);
        boundsPerEntity = new(32, Allocator.Persistent);
        this.ecsId = ecsId;
        this.chunkSize = chunkSize;
    }

    public void AddOrUpdate(Entity entity, AABB bounds, EntityTransform transform)
    {
        Remove(entity);
        var transformedBounds = transform * bounds;
        var minIndex = Index.PositionToIndex(transformedBounds.Min, chunkSize);
        var maxIndex = Index.PositionToIndex(transformedBounds.Max, chunkSize);

        for (int x = minIndex.x; x <= maxIndex.x; x++)
        {
            for (int z = minIndex.z; z <= maxIndex.z; z++)
            {
                var chunk = new int2(x, z);
                EnsureChunk(chunk).Add(entity.id);
                chunksPerEntity.Add(entity.id, chunk);
            }
        }
        boundsPerEntity[entity.id] = transformedBounds;
    }

    private UnsafeHashSetHolder<uint> EnsureChunk(int2 index)
    {
        if (!entitiesPerChunk.TryGetValue(index, out var entities))
        {
            entitiesPerChunk[index] = entities = new(32, Allocator.Persistent);
        }
        return entities;
    }

    public void Remove(Entity entity)
    {
        if (chunksPerEntity.TryGetValue(entity.id, out var chunks))
        {
            foreach (var chunk in chunks)
            {
                EnsureChunk(chunk).Remove(entity.id);
            }
            chunksPerEntity.Remove(entity.id);
            boundsPerEntity.Remove(entity.id);
        }
    }

    public NativeList<Entity> Query(AABB bounds, Allocator allocator)
    {
        var results = new NativeList<Entity>(128, allocator.AsJobSafe());
        var job = new QueryJob()
        {
            hashGrid = this,
            chunkSize = chunkSize,
            results = results
        }
        job.RunSafe();
        return results;
    }

    public void Dispose()
    {
        foreach (var item in entitiesPerChunk)
        {
            item.Value.Dispose();
        }
        entitiesPerChunk.Dispose();
        chunksPerEntity.Dispose();
        boundsPerEntity.Dispose();
    }

    [BurstCompile]
    public struct QueryJob : IJob
    {
        [ReadOnly] public EntityHashGrid hashGrid;
        public float chunkSize;
        public AABB bounds;
        public NativeList<Entity> results;

        public void Execute()
        {
            var minIndex = Index.PositionToIndex(transformedBounds.Min, chunkSize);
            var maxIndex = Index.PositionToIndex(transformedBounds.Max, chunkSize);

            var hashSet = new NativeHashSet<uint>(256, Allocator.Temp);
            for (int x = minIndex.x; x <= maxIndex.x; x++)
            {
                for (int z = minIndex.z; z <= maxIndex.z; z++)
                {
                    var chunk = new int2(x, z);
                    if (!hashGrid.entitiesPerChunk.TryGetValue(chunk, out var entities))
                        continue;
                    foreach (var entityId in entities)
                    {
                        if (hashSet.Add(entityId) && hashGrid.boundsPerEntity[entityId].Intersects(bounds))
                        {
                            results.Add(new Entity(hashGrid.ecsId, entityId));
                        }
                    }
                }
            }
        }
    }
}