using System.Collections;

namespace PersistentOrderedMap;

public sealed class PersistentOrderedMap<K, V, TStrategy> : BaseOrderedMap<K, V, TStrategy>, IEnumerable, IEnumerable<KeyValuePair<K, V>> where TStrategy : IKeyStrategy<K>
{
    internal PersistentOrderedMap(Node<K> root, TStrategy strategy, int count) 
        : base(root, strategy, count) { }

    // ---------------------------------------------------------
    // Immutable Write API (Returns new Map)
    // ---------------------------------------------------------
    public PersistentOrderedMap<K, V, TStrategy> Set(K key, V value)
    {
        // OPTIMIZATION: Use OwnerId.None (0).
        // This signals EnsureEditable to always copy the root path, 
        // producing a new tree of nodes that also have OwnerId.None.
        var newRoot = BTreeFunctions.Set(_root, key, value, _strategy, OwnerId.None, out bool countChanged);
        return new PersistentOrderedMap<K, V, TStrategy>(newRoot, _strategy, countChanged ? Count + 1 : Count);
    }
   
    public static PersistentOrderedMap<K, V, TStrategy> Empty(TStrategy strategy)
    {
        // Create an empty Leaf Node.
        // 'default(OwnerId)' (usually 0) marks this node as Immutable/Persistent.
        // This ensures that any subsequent Set/Remove will clone this node 
        // instead of modifying it in place.
        var emptyRoot = new LeafNode<K, V>(default(OwnerId), strategy.UsesPrefixes);
    
        return new PersistentOrderedMap<K, V, TStrategy>(emptyRoot, strategy, 0);
    }

    public PersistentOrderedMap<K, V, TStrategy> Remove(K key)
    {
        var newRoot = BTreeFunctions.Remove<K,V, TStrategy>(_root, key, _strategy, OwnerId.None, out bool removed);
        if (!removed) return this;
        return new PersistentOrderedMap<K, V, TStrategy>(newRoot, _strategy, Count - 1);
    }

    public TransientOrderedMap<K, V, TStrategy> ToTransient()
    {
        return new TransientOrderedMap<K, V, TStrategy>(_root, _strategy, Count);
    }
}