GhostEngine/Ghost.Entities/System.cs

using Ghost.Core;

namespace Ghost.Entities;

public readonly ref struct SystemAPI
{
    public TimeData Time
    {
        get; init;
    }
}

public interface ISystem
{
    World World
    {
        get; init;
    }

    void Initialize(ref readonly SystemAPI systemAPI);
    void Update(ref readonly SystemAPI systemAPI);
    void Cleanup(ref readonly SystemAPI systemAPI);
}

public abstract class SystemBase : ISystem
{
    private List<int>? _requiredQueries;

    public World World
    {
        get; init;
    } = null!;

    public int LastSystemVersion
    {
        get; internal set;
    } = -2;

    private bool ShouldUpdate()
    {
        if (_requiredQueries == null || _requiredQueries.Count == 0)
        {
            return true;
        }

        foreach (var queryID in _requiredQueries)
        {
            ref var query = ref World.GetEntityQueryReference(new Identifier<EntityQuery>(queryID));
            if (query.GetEntityCount() == 0)
            {
                return false;
            }
        }

        return true;
    }

    protected void RequireQueryForUpdate(Identifier<EntityQuery> queryID)
    {
        _requiredQueries ??= new List<int>(4);
        _requiredQueries.Add(queryID.Value);
    }

    void ISystem.Initialize(ref readonly SystemAPI systemAPI)
    {
        OnInitialize(in systemAPI);
    }

    void ISystem.Update(ref readonly SystemAPI systemAPI)
    {
        if (ShouldUpdate())
        {
            if (World.Version - LastSystemVersion > 1)
            {
                OnStartRunning();
            }

            OnUpdate(in systemAPI);
            LastSystemVersion = World.Version;
        }
        else
        {
            if (World.Version - LastSystemVersion <= 1)
            {
                OnStopRunning();
            }
        }
    }

    void ISystem.Cleanup(ref readonly SystemAPI systemAPI)
    {
        OnCleanup(in systemAPI);
    }

    protected virtual void OnInitialize(ref readonly SystemAPI systemAPI)
    {
    }

    protected virtual void OnUpdate(ref readonly SystemAPI systemAPI)
    {
    }

    protected virtual void OnCleanup(ref readonly SystemAPI systemAPI)
    {
    }

    protected virtual void OnStopRunning()
    {
    }

    protected virtual void OnStartRunning()
    {
    }
}

[AttributeUsage(AttributeTargets.Class | AttributeTargets.Struct, AllowMultiple = true)]
public class UpdateAfterAttribute : Attribute
{
    public Type SystemType { get; }

    public UpdateAfterAttribute(Type systemType)
    {
        SystemType = systemType;
    }
}

[AttributeUsage(AttributeTargets.Class | AttributeTargets.Struct, AllowMultiple = true)]
public class UpdateBeforeAttribute : Attribute
{
    public Type SystemType { get; }

    public UpdateBeforeAttribute(Type systemType)
    {
        SystemType = systemType;
    }
}

[AttributeUsage(AttributeTargets.Class | AttributeTargets.Struct, AllowMultiple = false)]
public class SystemGroupAttribute : Attribute
{
    public Type GroupType { get; }

    public SystemGroupAttribute(Type groupType)
    {
        GroupType = groupType;
    }
}

#if false
internal static partial class SystemGroupRegistry
{
    private static readonly Dictionary<Type, List<ISystem>> _systemGroupMap = new();

    // TODO: Use Source Generators to generate group registrations at compile time.

    public static void RegisterSystemGroup<T>(Type groupType)
        where T : ISystem, new()
    {
        if (!_systemGroupMap.ContainsKey(typeof(T)))
        {
            _systemGroupMap[typeof(T)] = new();
        }

        _systemGroupMap[groupType].Add(new T());
    }

    public static List<ISystem> GetSystemsForGroup(Type groupType)
    {
        if (_systemGroupMap.TryGetValue(groupType, out var systems))
        {
            return systems;
        }

        throw new InvalidOperationException($"No systems registered for System Group of type {groupType.FullName}");
    }
}
#endif

public abstract class SystemGroup : ISystem
{
    private readonly List<ISystem> _systems = [];
    private List<ISystem>? _sortedSystems;

    private uint _version = 0;
    private uint _sortedVersion = 0;

    public World World
    {
        get; init;
    } = null!;

    // public SystemGroup()
    // {
    //     _systems = SystemGroupRegistry.GetSystemsForGroup(GetType());
    // }

    private static List<ISystem> Sort(List<ISystem> systems)
    {
        // 1. Build the Graph
        // Key: The System, Value: Systems that MUST run before the Key
        var dependencies = new Dictionary<Type, HashSet<Type>>();
        var systemMap = systems.ToDictionary(s => s.GetType(), s => s);

        foreach (var sys in systems)
        {
            var type = sys.GetType();
            if (!dependencies.TryGetValue(type, out HashSet<Type>? value))
            {
                value = [];
                dependencies[type] = value;
            }

            // Handle [UpdateAfter(typeof(Other))] -> Other comes before This
            foreach (var attr in type.GetCustomAttributes(typeof(UpdateAfterAttribute), true))
            {
                var depType = ((UpdateAfterAttribute)attr).SystemType;
                value.Add(depType);
            }

            // Handle [UpdateBefore(typeof(Other))] -> This comes before Other
            // Which means: Other depends on This
            foreach (var attr in type.GetCustomAttributes(typeof(UpdateBeforeAttribute), true))
            {
                var targetType = ((UpdateBeforeAttribute)attr).SystemType;
                if (!dependencies.ContainsKey(targetType)) dependencies[targetType] = [];
                dependencies[targetType].Add(type);
            }
        }

        // 2. Topological Sort (Kahn's Algorithm variant)
        var sortedList = new List<ISystem>();
        var visited = new HashSet<Type>();

        // We loop until we have sorted everyone
        while (sortedList.Count < systems.Count)
        {
            bool addedAny = false;

            foreach (var sys in systems)
            {
                var type = sys.GetType();
                if (visited.Contains(type)) continue;

                // Check if all dependencies for this system are already visited/sorted
                bool canRun = true;
                if (dependencies.TryGetValue(type, out var deps))
                {
                    foreach (var dep in deps)
                    {
                        // If the dependency exists in our list but hasn't run yet, we can't run.
                        // (We check systemMap to ignore dependencies that don't exist in this world)
                        if (systemMap.ContainsKey(dep) && !visited.Contains(dep))
                        {
                            canRun = false;
                            break;
                        }
                    }
                }

                if (canRun)
                {
                    sortedList.Add(sys);
                    visited.Add(type);
                    addedAny = true;
                }
            }

            if (!addedAny)
            {
                throw new InvalidOperationException("Circular Dependency detected in Systems! Check your [UpdateAfter] attributes.");
            }
        }

        return sortedList;
    }

    public void AddSystem<T>()
        where T : ISystem, new()
    {
        _systems.Add(new T()
        {
            World = World
        });

        _version++;
    }

    public void SortSystems()
    {
        if (_sortedVersion == _version)
        {
            return;
        }

        if (_systems.Count == 0)
        {
            _sortedSystems = [];
            _sortedVersion = _version;
            return;
        }

        _sortedSystems = Sort(_systems);
        _sortedVersion = _version;
    }

    private void ThrowIfNotSorted()
    {
        if (_sortedSystems == null || _sortedVersion != _version)
        {
            throw new InvalidOperationException("Systems must be sorted before calling this method. Call SortSystems() after adding all systems.");
        }
    }

    public void Initialize(ref readonly SystemAPI systemAPI)
    {
        ThrowIfNotSorted();

        foreach (var system in _sortedSystems!)
        {
            system.Initialize(in systemAPI);
        }
    }

    public void Update(ref readonly SystemAPI systemAPI)
    {
        ThrowIfNotSorted();

        foreach (var system in _sortedSystems!)
        {
            system.Update(in systemAPI);
        }
    }

    public void Cleanup(ref readonly SystemAPI systemAPI)
    {
        ThrowIfNotSorted();

        foreach (var system in _sortedSystems!)
        {
            system.Cleanup(in systemAPI);
        }
    }
}

public class DefaultSystemGroup : SystemGroup
{
}

public class SystemManager
{
    private readonly World _world;

    private readonly List<ISystem> _systems = [];

    internal IReadOnlyList<ISystem> Systems => _systems;

    internal SystemManager(World world)
    {
        _world = world;
        AddSystem<DefaultSystemGroup>();
    }

    public void AddSystem<T>()
        where T : ISystem, new()
    {
        _systems.Add(new T()
        {
            World = _world
        });
    }

    public T GetSystem<T>()
        where T : ISystem
    {
        foreach (var system in _systems)
        {
            if (system is T typedSystem)
            {
                return typedSystem;
            }
        }

        throw new InvalidOperationException($"System of type {typeof(T).FullName} not found in SystemManager.");
    }

    internal void InitializeAll(ref readonly SystemAPI systemAPI)
    {
        foreach (var system in _systems)
        {
            system.Initialize(in systemAPI);
        }
    }

    internal void UpdateAll(ref readonly SystemAPI systemAPI)
    {
        foreach (var system in _systems)
        {
            system.Update(in systemAPI);
        }
    }

    internal void CleanupAll(ref readonly SystemAPI systemAPI)
    {
        foreach (var system in _systems)
        {
            system.Cleanup(in systemAPI);
        }
    }
}