using System;
using System.Collections.Generic;

namespace Engine.Core;

public class BehaviourCollectorOrdered<TIndex, TItem> : BehaviourCollectorBase<TItem> where TItem : class where TIndex : IComparable
{
    private readonly SortedDictionary<TIndex, FastList<TItem>> behaviours = null!;
    private readonly Dictionary<TItem, TIndex> indexCache = null!;

    private readonly Func<TItem, TIndex> getIndexFunc = null!;
    private readonly IComparer<TIndex> sortBy = null!;

    private int count = 0;
    public override int Count => count;

    public override TItem this[Index index]
    {
        get
        {
            int actualIndex = index.IsFromEnd
                ? count - index.Value
                : index.Value;

            if (actualIndex < 0 || actualIndex >= count)
                throw new IndexOutOfRangeException();

            int leftIndex = actualIndex;
            foreach ((TIndex i, FastList<TItem> list) in behaviours)
            {
                if (leftIndex < list.Count)
                    return list[leftIndex];
                leftIndex -= list.Count;
            }
            throw new IndexOutOfRangeException();
        }
    }

    protected override bool RemoveBehaviour(TItem tBehaviour)
    {
        if (!indexCache.TryGetValue(tBehaviour, out TIndex? index) || !behaviours.TryGetValue(index, out FastList<TItem>? list))
            throw new Exceptions.NotFoundException($"Index of '{index}' is not found in the collector");

        if (!list.Remove(tBehaviour) || !indexCache.Remove(tBehaviour))
            return false;

        count--;
        return true;
    }

    protected override void AddBehaviour(TItem behaviour)
    {
        TIndex key = getIndexFunc(behaviour);
        if (!behaviours.TryGetValue(key, out FastList<TItem>? list))
            behaviours[key] = list = [];

        count++;
        list.Add(behaviour);
        indexCache.Add(behaviour, key);
    }

    public BehaviourCollectorOrdered(Func<TItem, TIndex> getIndexFunc, Comparison<TIndex> sortBy)
    {
        this.getIndexFunc = getIndexFunc;
        this.sortBy = Comparer<TIndex>.Create(sortBy);
        behaviours = new(this.sortBy);
    }

    public BehaviourCollectorOrdered(IUniverse universe, Func<TItem, TIndex> getIndexFunc, Comparison<TIndex> sortBy) : base(universe)
    {
        this.getIndexFunc = getIndexFunc;
        this.sortBy = Comparer<TIndex>.Create(sortBy);
        behaviours = new(this.sortBy);
    }

    public BehaviourCollectorOrdered(Func<TItem, TIndex> getIndexFunc, IComparer<TIndex> sortBy)
    {
        this.getIndexFunc = getIndexFunc;
        this.sortBy = sortBy;
        behaviours = new(sortBy);
    }

    public BehaviourCollectorOrdered(IUniverse universe, Func<TItem, TIndex> getIndexFunc, IComparer<TIndex> sortBy) : base(universe)
    {
        this.getIndexFunc = getIndexFunc;
        this.sortBy = sortBy;
        behaviours = new(sortBy);
    }
}