using R3;

namespace OECS;

/// <summary>
/// Buffers changes during system execution and posts them to R3 subjects
/// when <see cref="Post"/> is called.
///
/// Subscribers receive batched changes: all changes accumulated since the
/// last <see cref="Post"/> call are pushed in a single burst.
/// </summary>
internal class ChangeBuffer
{
    private readonly ChangeSet _pending = new();

    private readonly Subject<EntityChange> _entitySubject = new();
    private readonly Dictionary<Type, TrackedSubject> _componentSubjects = new();
    private readonly Dictionary<QueryDescriptor, TrackedSubject> _querySubjects = new();

    /// <summary>
    /// Wraps a <see cref="Subject{T}"/> with a subscriber count so that
    /// subjects with no remaining subscribers can be cleaned up.
    /// </summary>
    private sealed class TrackedSubject
    {
        public readonly Subject<EntityChange> Subject = new();
        public int SubscriberCount;
    }

    /// <summary>
    /// The change set currently accumulating. Cleared after each <see cref="Post"/>.
    /// </summary>
    public ChangeSet Pending => _pending;

    /// <summary>
    /// Pushes all pending changes to R3 subjects, then clears the pending set.
    /// </summary>
    public void Post()
    {
        if (_pending.Count == 0)
            return;

        var changes = _pending.Changes;

        foreach (var change in changes)
        {
            // Push to the global entity subject.
            _entitySubject.OnNext(change);

            // Push to component-specific subjects.
            if (change.ComponentType != null)
            {
                if (_componentSubjects.TryGetValue(change.ComponentType, out var compTracked))
                {
                    compTracked.Subject.OnNext(change);
                }
            }

            // Push to matching query subjects.
            foreach (var (query, tracked) in _querySubjects)
            {
                if (ChangeMatchesQuery(change, query))
                {
                    tracked.Subject.OnNext(change);
                }
            }
        }

        _pending.Clear();
    }

    /// <summary>
    /// Returns an observable that emits all entity changes.
    /// </summary>
    public Observable<EntityChange> ObserveEntityChanges()
    {
        return _entitySubject;
    }

    /// <summary>
    /// Returns an observable that emits changes for a specific component type.
    /// </summary>
    public Observable<EntityChange> ObserveComponentChanges(Type componentType)
    {
        if (!_componentSubjects.TryGetValue(componentType, out var tracked))
        {
            tracked = new TrackedSubject();
            _componentSubjects[componentType] = tracked;
        }

        tracked.SubscriberCount++;
        return WrapWithCleanup(tracked.Subject, () =>
        {
            tracked.SubscriberCount--;
            if (tracked.SubscriberCount <= 0)
            {
                tracked.Subject.Dispose();
                _componentSubjects.Remove(componentType);
            }
        });
    }

    /// <summary>
    /// Returns an observable that emits changes matching the given query.
    /// </summary>
    public Observable<EntityChange> ObserveQuery(QueryDescriptor query)
    {
        if (!_querySubjects.TryGetValue(query, out var tracked))
        {
            tracked = new TrackedSubject();
            _querySubjects[query] = tracked;
        }

        tracked.SubscriberCount++;
        return WrapWithCleanup(tracked.Subject, () =>
        {
            tracked.SubscriberCount--;
            if (tracked.SubscriberCount <= 0)
            {
                tracked.Subject.Dispose();
                _querySubjects.Remove(query);
            }
        });
    }

    /// <summary>
    /// Wraps an observable so that <paramref name="onLastDispose"/> is called
    /// when the last subscriber disposes.
    /// </summary>
    private static Observable<EntityChange> WrapWithCleanup(
        Observable<EntityChange> source, Action onLastDispose)
    {
        return Observable.Create<EntityChange>(observer =>
        {
            var subscription = source.Subscribe(observer);
            return Disposable.Create(() =>
            {
                subscription.Dispose();
                onLastDispose();
            });
        });
    }

    /// <summary>
    /// Disposes all R3 subjects.
    /// </summary>
    public void Dispose()
    {
        _entitySubject.Dispose();
        foreach (var tracked in _componentSubjects.Values)
        {
            tracked.Subject.Dispose();
        }
        _componentSubjects.Clear();
        foreach (var tracked in _querySubjects.Values)
        {
            tracked.Subject.Dispose();
        }
        _querySubjects.Clear();
    }

    private static bool ChangeMatchesQuery(EntityChange change, QueryDescriptor query)
    {
        // Entity-level changes: match if the entity was added/removed and
        // the query has any "with" types (we can't know component state for
        // removed entities, so we only match added entities that would satisfy
        // the query — but we don't have component data here).
        // For simplicity, we only match component-level changes against queries.
        if (change.ComponentType == null)
            return false;

        // A component change matches a query if the component type is in the
        // query's With set and not in the Without set.
        if (query.Without.Contains(change.ComponentType))
            return false;

        return query.With.Contains(change.ComponentType);
    }
}