using System.Diagnostics;
using SJK.Functional;

public class Palette<TEntry> where TEntry : IEquatable<TEntry>

{
    private List<TEntry> _entries = new();
    private Dictionary<TEntry, int> _index = new();
    public int Capacity => _entries.Capacity;
    public int Count => _entries.Count;
    public TEntry Get(int index)
    {
        Debug.Assert(index >= 0 && index < Count,$"{nameof(index)} is out of range, {nameof(index)} has value {index}, while {nameof(Count)} has value of {Count}");
        return _entries[index];
    }
    public Option<TEntry> Get(Predicate<TEntry> predicate)
    {
        foreach (var item in _entries)
        {
            if (predicate(item))
            {
                return Option<TEntry>.Some(item);
            }
        }
        return Option<TEntry>.None;
    }
    public virtual int GetOrAddIndex(TEntry entry)
    {
        int index;
        if (_index.TryGetValue(entry, out index))
            return index;
        index = _entries.Count;
        _index.Add(entry, index);
        _entries.Add(entry);
        return index;
    }
    
    public IReadOnlyList<TEntry> Entries => _entries;
    public void Clear()
    {
        _entries.Clear();
        _index.Clear();
    }
}


public class RefPalette<TEntry> where TEntry : IEquatable<TEntry>
{
    private readonly List<TEntry> _entries = new();
    private readonly List<int> _refCounts = new();
    private readonly Dictionary<TEntry, int> _index = new();
    private readonly Stack<int> _free = new();

    public RefPalette(TEntry defaultEntry)
    {
        _entries.Add(defaultEntry);
        _refCounts.Add(int.MaxValue); // default entry never removed
        _index.Add(defaultEntry, 0);
    }

    public int Count => _entries.Count;
    public IReadOnlyList<TEntry> Entries => _entries;

    public TEntry Get(int index)
    {
        Debug.Assert(index >= 0 && index < _entries.Count,
            $"{nameof(index)} out of range. index={index}, count={_entries.Count}");

        return _entries[index];
    }

    public int GetOrAddIndex(TEntry entry)
    {
        if (_index.TryGetValue(entry, out var index))
            return index;

        if (_free.Count > 0)
        {
            index = _free.Pop();
            _entries[index] = entry;
            _refCounts[index] = 0;
        }
        else
        {
            index = _entries.Count;
            _entries.Add(entry);
            _refCounts.Add(0);
        }

        _index[entry] = index;
        return index;
    }

    public void AddRef(int index)
    {
        Debug.Assert(index >= 0 && index < _refCounts.Count);
        _refCounts[index]++;
    }

    public void Release(int index)
    {
        Debug.Assert(index > 0 && index < _refCounts.Count); // never release default

        int count = --_refCounts[index];

        if (count == 0)
        {
            var entry = _entries[index];
            _index.Remove(entry);

            _entries[index] = default!;
            _free.Push(index);
        }
    }

    public int RefCount(int index)
    {
        return _refCounts[index];
    }

    public Option<TEntry> Get(Predicate<TEntry> predicate)
    {
        for (int i = 0; i < _entries.Count; i++)
        {
            if (_refCounts[i] > 0 && predicate(_entries[i]))
                return Option<TEntry>.Some(_entries[i]);
        }

        return Option<TEntry>.None;
    }
}