using System.Runtime.CompilerServices;

namespace OECS;

/// <summary>
/// Maintains a reverse index for relationship components, mapping
/// target entities to the set of source entities that point to them.
///
/// Updated automatically by <see cref="World"/> when relationship
/// components are added, removed, or when entities are destroyed.
/// </summary>
internal class RelationshipIndex
{
    /// <summary>
    /// Per relationship type: target entity → set of source entities.
    /// </summary>
    private readonly Dictionary<Type, Dictionary<Entity, HashSet<Entity>>> _index = new();

    /// <summary>
    /// Registers that a source entity now has a relationship pointing to a target.
    /// </summary>
    public void OnAdded(Type relationshipType, Entity source, Entity target)
    {
        if (!_index.TryGetValue(relationshipType, out var targetMap))
        {
            targetMap = new Dictionary<Entity, HashSet<Entity>>();
            _index[relationshipType] = targetMap;
        }

        if (!targetMap.TryGetValue(target, out var sources))
        {
            sources = new HashSet<Entity>();
            targetMap[target] = sources;
        }

        sources.Add(source);
    }

    /// <summary>
    /// Unregisters a relationship from a source entity pointing to a target.
    /// </summary>
    public void OnRemoved(Type relationshipType, Entity source, Entity target)
    {
        if (!_index.TryGetValue(relationshipType, out var targetMap))
            return;

        if (!targetMap.TryGetValue(target, out var sources))
            return;

        sources.Remove(source);

        if (sources.Count == 0)
            targetMap.Remove(target);
    }

    /// <summary>
    /// Returns all source entities that have a relationship of the given type
    /// pointing to the specified target entity.
    /// </summary>
    public IReadOnlyCollection<Entity> GetSources<T>(Entity target)
        where T : struct, IRelationship
    {
        var type = typeof(T);
        if (!_index.TryGetValue(type, out var targetMap))
            return Array.Empty<Entity>();

        if (!targetMap.TryGetValue(target, out var sources))
            return Array.Empty<Entity>();

        return sources;
    }

    /// <summary>
    /// Returns all target entities that the given source entity points to
    /// via the specified relationship type. Returns an empty collection if none.
    /// </summary>
    public IReadOnlyCollection<Entity> GetTargets(Type relationshipType, Entity source)
    {
        if (!_index.TryGetValue(relationshipType, out var targetMap))
            return Array.Empty<Entity>();

        var result = new List<Entity>();
        foreach (var (target, sources) in targetMap)
        {
            if (sources.Contains(source))
                result.Add(target);
        }
        return result;
    }

    /// <summary>
    /// Removes all index entries where the given entity appears as a source.
    /// Called before the entity's components are removed during destruction.
    /// </summary>
    public void RemoveAllSourcesForEntity(Entity entity)
    {
        foreach (var (_, targetMap) in _index)
        {
            foreach (var (target, sources) in targetMap)
            {
                sources.Remove(entity);
            }
        }

        // Clean up empty target entries.
        foreach (var (_, targetMap) in _index)
        {
            var emptyTargets = new List<Entity>();
            foreach (var (target, sources) in targetMap)
            {
                if (sources.Count == 0)
                    emptyTargets.Add(target);
            }
            foreach (var target in emptyTargets)
            {
                targetMap.Remove(target);
            }
        }
    }

    /// <summary>
    /// Returns all relationship types for which the given entity is a target,
    /// along with the set of source entities pointing to it.
    /// Used during entity destruction to cascade-remove incoming relationships.
    /// </summary>
    public IEnumerable<(Type RelationshipType, IReadOnlyCollection<Entity> Sources)> GetIncomingRelationships(Entity target)
    {
        foreach (var (relType, targetMap) in _index)
        {
            if (targetMap.TryGetValue(target, out var sources) && sources.Count > 0)
            {
                yield return (relType, sources);
            }
        }
    }

    /// <summary>
    /// Returns true if the index has any entries for the given relationship type.
    /// </summary>
    public bool HasType(Type relationshipType)
    {
        return _index.ContainsKey(relationshipType);
    }
}