using System;
using System.Collections.Generic;

namespace Encompass
{
    internal class TrackingManager
    {
        private readonly Dictionary<Type, HashSet<Engine>> _immediateComponentTypesToEngines = new Dictionary<Type, HashSet<Engine>>();
        private readonly Dictionary<Type, HashSet<Engine>> _componentTypesToEngines = new Dictionary<Type, HashSet<Engine>>();

        private readonly HashSet<(int, Type)> _additions = new HashSet<(int, Type)>();
        private readonly HashSet<(int, Type)> _removals = new HashSet<(int, Type)>();

        private readonly HashSet<(int, Engine)> _pairsToCheck = new HashSet<(int, Engine)>();

        public void RegisterComponentTypeToEngine(Type type, Engine engine)
        {
            if (!_componentTypesToEngines.ContainsKey(type)) { _componentTypesToEngines.Add(type, new HashSet<Engine>()); }
            _componentTypesToEngines[type].Add(engine);
        }

        public void RegisterImmediateComponentTypeToEngine(Type type, Engine engine)
        {
            if (!_immediateComponentTypesToEngines.ContainsKey(type)) { _immediateComponentTypesToEngines.Add(type, new HashSet<Engine>()); }
            _immediateComponentTypesToEngines[type].Add(engine);
        }

        public void RegisterAddition(int entityID, Type type)
        {
            _additions.Add((entityID, type));
        }

        public void RegisterRemoval(int entityID, Type type)
        {
            _removals.Add((entityID, type));
        }

        public void InitializeTracking(IEnumerable<int> entityIDs)
        {
            foreach (var entityID in entityIDs)
            {
                foreach (var engineSet in _componentTypesToEngines.Values)
                {
                    foreach (var engine in engineSet)
                    {
                        engine.CheckAndUpdateTracking(entityID);
                    }
                }
            }
        }

        public void ImmediateUpdateTracking(int entityID, Type componentType)
        {
            if (_immediateComponentTypesToEngines.ContainsKey(componentType))
            {
                foreach (var engine in _componentTypesToEngines[componentType])
                {
                    engine.ImmediateCheckAndUpdateTracking(entityID);
                }
            }
        }

        public void UpdateTracking()
        {
            // TODO: optimize so we only check each entity/engine pair once
            foreach (var (entity, componentType) in _additions)
            {
                if (_componentTypesToEngines.ContainsKey(componentType))
                {
                    foreach (var engine in _componentTypesToEngines[componentType])
                    {
                        _pairsToCheck.Add((entity, engine));
                    }
                }
            }
            _additions.Clear();

            foreach (var (entity, componentType) in _removals)
            {
                if (_componentTypesToEngines.ContainsKey(componentType))
                {
                    foreach (var engine in _componentTypesToEngines[componentType])
                    {
                        _pairsToCheck.Add((entity, engine));
                    }
                }
            }
            _removals.Clear();

            foreach (var (entity, engine) in _pairsToCheck)
            {
                engine.CheckAndUpdateTracking(entity);
            }
            _pairsToCheck.Clear();
        }
    }
}