using System.Runtime.CompilerServices;

namespace OECS;

/// <summary>
/// Registry of all component sparse sets, keyed by component type.
/// Provides typed generic accessors that delegate to the underlying
/// <see cref="SparseSet{T}"/> instances.
/// </summary>
internal class ComponentStore
{
    private readonly Dictionary<Type, ISparseSet> _sets = new();

    /// <summary>
    /// Gets or creates the sparse set for component type <typeparamref name="T"/>.
    /// </summary>
    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    private SparseSet<T> GetSet<T>() where T : struct
    {
        var type = typeof(T);
        if (!_sets.TryGetValue(type, out var set))
        {
            var newSet = new SparseSet<T>();
            _sets[type] = newSet;
            return newSet;
        }
        return (SparseSet<T>)set;
    }

    /// <summary>
    /// Gets the sparse set for component type <typeparamref name="T"/> if it exists.
    /// </summary>
    private SparseSet<T>? GetSetIfExists<T>() where T : struct
    {
        if (_sets.TryGetValue(typeof(T), out var set))
            return (SparseSet<T>)set;
        return null;
    }

    /// <summary>
    /// Returns all registered component types.
    /// </summary>
    public IEnumerable<Type> ComponentTypes => _sets.Keys;

    public void Add<T>(Entity entity, T component) where T : struct
    {
        GetSet<T>().Add(entity, component);
    }

    public void Remove<T>(Entity entity) where T : struct
    {
        GetSetIfExists<T>()?.Remove(entity);
    }

    public ref T Get<T>(Entity entity) where T : struct
    {
        return ref GetSet<T>().Get(entity);
    }

    public bool Has<T>(Entity entity) where T : struct
    {
        return GetSetIfExists<T>()?.Contains(entity) ?? false;
    }

    /// <summary>
    /// Removes all components for the given entity across all component types.
    /// </summary>
    public void RemoveAll(Entity entity)
    {
        foreach (var set in _sets.Values)
        {
            set.Remove(entity);
        }
    }

    /// <summary>
    /// Returns the sparse set for a given type, or null if not registered.
    /// Used by query execution to probe sets by Type.
    /// </summary>
    public ISparseSet? GetSet(Type componentType)
    {
        _sets.TryGetValue(componentType, out var set);
        return set;
    }

    /// <summary>
    /// Returns the count of entities in the sparse set for type <typeparamref name="T"/>.
    /// Returns 0 if the set does not exist.
    /// </summary>
    public int Count<T>() where T : struct
    {
        return GetSetIfExists<T>()?.Count ?? 0;
    }

    /// <summary>
    /// Returns the count of entities in the sparse set for the given type.
    /// Returns 0 if the set does not exist.
    /// </summary>
    public int Count(Type componentType)
    {
        return GetSet(componentType)?.Count ?? 0;
    }
}