TechbloxModdingAPI/TechbloxModdingAPI/Utility/NativeApiExtensions.cs

using System;
using System.Collections.Generic;
using System.Reflection;
using HarmonyLib;
using Svelto.DataStructures;
using Svelto.ECS;
using Svelto.Tasks;
using Svelto.Tasks.Lean;
using TechbloxModdingAPI.Tasks;

namespace TechbloxModdingAPI.Utility
{
    public static class NativeApiExtensions
    {
        /// <summary>
        /// Attempts to query an entity and returns an optional that contains the result if succeeded.
        /// <b>This overload does not take initializer data into account.</b>
        /// </summary>
        /// <param name="entitiesDB">The entities DB</param>
        /// <param name="egid">The EGID to query</param>
        /// <typeparam name="T">The component type to query</typeparam>
        /// <returns>An optional that contains the result on success or is empty if not found</returns>
        public static OptionalRef<T> QueryEntityOptional<T>(this EntitiesDB entitiesDB, EGID egid)
            where T : unmanaged, IEntityComponent
        {
            return entitiesDB.TryQueryEntitiesAndIndex<T>(egid, out uint index, out var array)
                ? new OptionalRef<T>(array, index)
                : new OptionalRef<T>();
        }

        /// <summary>
        /// Attempts to query an entity and returns the result in an optional reference.
        /// </summary>
        /// <param name="entitiesDB">The entities DB to query from</param>
        /// <param name="obj">The ECS object to query</param>
        /// <param name="group">The group of the entity if the object can have multiple</param>
        /// <typeparam name="T">The component to query</typeparam>
        /// <returns>A reference to the component or a dummy value</returns>
        public static OptionalRef<T> QueryEntityOptional<T>(this EntitiesDB entitiesDB, EcsObjectBase obj,
            ExclusiveGroupStruct group = default)
            where T : unmanaged, IEntityComponent
        {
            EGID id = group == ExclusiveGroupStruct.Invalid ? obj.Id : new EGID(obj.Id.entityID, group);
            var opt = QueryEntityOptional<T>(entitiesDB, id);
            return opt ? opt : new OptionalRef<T>(obj, true);
        }

        /// <summary>
        /// Attempts to query an entity and returns the result or a dummy value that can be modified.
        /// </summary>
        /// <param name="entitiesDB">The entities DB to query from</param>
        /// <param name="obj">The ECS object to query</param>
        /// <param name="group">The group of the entity if the object can have multiple</param>
        /// <typeparam name="T">The component to query</typeparam>
        /// <returns>A reference to the component or a dummy value</returns>
        public static ref T QueryEntityOrDefault<T>(this EntitiesDB entitiesDB, EcsObjectBase obj,
            ExclusiveGroupStruct group = default)
            where T : unmanaged, IEntityComponent
        {
            EGID id = group == ExclusiveGroupStruct.Invalid ? obj.Id : new EGID(obj.Id.entityID, group);
            var opt = QueryEntityOptional<T>(entitiesDB, id);
            if (opt) return ref opt.Get();
            if (obj.InitData.Valid) return ref obj.InitData.Initializer(id).GetOrAdd<T>();
            /*if (!obj.InitData.Valid) return ref opt.Get(); //Default value
            var init = obj.InitData.Initializer(id);
            // Do not create the component if missing, as that can trigger Add() listeners that, in some cases, may be
            // invalid if (ab)using the classes in an unusual way - TODO: Check entity descriptor or something
            if (init.Has<T>()) return ref init.Get<T>();*/
            return ref opt.Get(); //Default value
        }

        private static readonly Dictionary<Type, (int PublishedCount, HashSet<EGID> Changes)> ChangesToPublish = new();

        /// <summary>
        /// Publishes an entity change, ignoring duplicate publishes and delaying changes as necessary.
        /// It will only publish in the next frame.
        /// </summary>
        /// <param name="entitiesDB">The entities DB to publish to</param>
        /// <param name="id">The ECS object that got changed</param>
        /// <param name="limit">Limits how many changes to publish - should be no more than the consumers' capacity that process this component</param>
        /// <typeparam name="T">The component that changed</typeparam>
        public static void PublishEntityChangeDelayed<T>(this EntitiesDB entitiesDB, EGID id, int limit = 80)
            where T : unmanaged, IEntityComponent
        {
            if (!ChangesToPublish.ContainsKey(typeof(T)))
                ChangesToPublish.Add(typeof(T), (0, new HashSet<EGID>()));
            var changes = ChangesToPublish[typeof(T)].Changes;
            if (changes.Contains(id)) return;
            changes.Add(id);
            PublishChanges<T>(entitiesDB, id, limit).RunOn(Scheduler.leanRunner);
        }

        private static IEnumerator<TaskContract> PublishChanges<T>(EntitiesDB entitiesDB, EGID id, int limit)
            where T : unmanaged, IEntityComponent
        {
            yield return Yield.It;
            while (ChangesToPublish[typeof(T)].PublishedCount >= limit)
                yield return Yield.It;
            if (!entitiesDB._entityStream._streams.TryGetValue(TypeRefWrapper<T>.wrapper, out var result))
                yield break; // There is no entity stream for this type
            var consumers = (result as EntityStream<T>)?._consumers;
            if (consumers == null)
            {
                Console.WriteLine("Consumers is null");
                yield break;
            }

            bool waitForConsumers;
            do
            {
                waitForConsumers = false;
                for (int i = 0; i < consumers.count; i++)
                {
                    var buffer = consumers[i]._ringBuffer;
                    if (buffer.Count + 1 <= buffer.Capacity) continue;
                    waitForConsumers = true;
                    Console.WriteLine($"Gonna have to wait for a consumer (capacity: {buffer.Capacity} count: {buffer.Count}");
                    break;
                }

                if (waitForConsumers) yield return Yield.It;
            } while (waitForConsumers);
            entitiesDB.PublishEntityChange<T>(id);
            var (count, changes) = ChangesToPublish[typeof(T)];
            changes.Remove(id);
            ChangesToPublish[typeof(T)] = (count + 1, changes);
            yield return Yield.It;
            ChangesToPublish[typeof(T)] = (Math.Max(ChangesToPublish[typeof(T)].PublishedCount - 1, 0), changes);
        }

        /// <summary>
        /// Query entities as OptionalRefs. The elements always exist, it's just a nice way to encapsulate the data.
        /// </summary>
        /// <param name="entitiesDB"></param>
        /// <param name="group"></param>
        /// <param name="select"></param>
        /// <typeparam name="T"></typeparam>
        /// <typeparam name="TR"></typeparam>
        /// <returns></returns>
        public static RefCollection<T> QueryEntitiesOptional<T>(this EntitiesDB entitiesDB, ExclusiveGroupStruct group) where T : unmanaged, IEntityComponent
        {
            var (buffer, ids, count) = entitiesDB.QueryEntities<T>(group);
            return new RefCollection<T>(count, buffer, ids, group);
        }
    }
}