using System.Collections;

namespace PersistentOrderedMap;

public abstract class BaseOrderedMap<K, V, TStrategy> : IEnumerable<KeyValuePair<K, V>> where TStrategy : IKeyStrategy<K>
{
    internal Node<K> _root;
    internal readonly TStrategy _strategy;
    
    public int Count { get; protected set; }

    protected BaseOrderedMap(Node<K> root, TStrategy strategy, int count)
    {
        _root = root ?? throw new ArgumentNullException(nameof(root));
        _strategy = strategy ?? throw new ArgumentNullException(nameof(strategy));
        Count = count;

    }

    // ---------------------------------------------------------
    // Read Operations (Shared)
    // ---------------------------------------------------------

    public bool TryGetValue(K key, out V value)
    {
        return BTreeFunctions.TryGetValue(_root, key, _strategy, out value);
    }

    public bool ContainsKey(K key)
    {
        return BTreeFunctions.TryGetValue<K,V, TStrategy>(_root, key, _strategy, out _);
    }
    
    

    // ---------------------------------------------------------
    // Bootstrap / Factory Helpers
    // ---------------------------------------------------------
        
    public static PersistentOrderedMap<K, V, TStrategy> Create(TStrategy strategy)
    {
        // Start with an empty leaf owned by None so the first write triggers CoW.
        var emptyRoot = new LeafNode<K, V>(OwnerId.None, strategy.UsesPrefixes);
        return new PersistentOrderedMap<K, V, TStrategy>(emptyRoot, strategy, 0);
    }

    public static TransientOrderedMap<K, V, TStrategy> CreateTransient(TStrategy strategy)
    {
        var emptyRoot = new LeafNode<K, V>(OwnerId.None, strategy.UsesPrefixes);
        return new TransientOrderedMap<K, V, TStrategy>(emptyRoot, strategy,0);
    }
    
    
    public BTreeEnumerator<K, V, TStrategy> GetEnumerator()
    {
        return AsEnumerable().GetEnumerator();
    }

    IEnumerator<KeyValuePair<K, V>> IEnumerable<KeyValuePair<K, V>>.GetEnumerator()
    {
        return GetEnumerator();
    }

    IEnumerator IEnumerable.GetEnumerator()
    {
        return GetEnumerator();
    }
    
    // 1. Full Scan
    public BTreeEnumerable<K, V, TStrategy> AsEnumerable() 
        => new(_root, _strategy, false, default, false, default);

// 2. Exact Range
    public BTreeEnumerable<K, V, TStrategy> Range(K min, K max) 
        => new(_root, _strategy, true, min, true, max);

// 3. Start From (Open Ended)
    public BTreeEnumerable<K, V, TStrategy> From(K min) => new(_root, _strategy, true, min, false, default);

// 4. Until (Start at beginning)
    public BTreeEnumerable<K, V, TStrategy> Until(K max) 
        => new(_root, _strategy, false, default, true, max);



    // ---------------------------------------------------------
// Navigation Operations
// ---------------------------------------------------------

public bool TryGetMin(out K key, out V value) => BTreeFunctions.TryGetMin(_root, out key, out value);

public bool TryGetMax(out K key, out V value) => BTreeFunctions.TryGetMax(_root, out key, out value);

public bool TryGetSuccessor(K key, out K nextKey, out V nextValue) => BTreeFunctions.TryGetSuccessor(_root, key, _strategy, out nextKey, out nextValue);

public bool TryGetPredecessor(K key, out K prevKey, out V prevValue) => BTreeFunctions.TryGetPredecessor(_root, key, _strategy, out prevKey, out prevValue);

// ---------------------------------------------------------
// Set Operations (Linear Merge O(N+M))
// ---------------------------------------------------------

public IEnumerable<KeyValuePair<K, V>> Intersect(BaseOrderedMap<K, V, TStrategy> other)
{
    using var enum1 = this.GetEnumerator();
    using var enum2 = other.GetEnumerator();
    
    bool has1 = enum1.MoveNext();
    bool has2 = enum2.MoveNext();

    while (has1 && has2)
    {
        int cmp = _strategy.Compare(enum1.Current.Key, enum2.Current.Key);
        if (cmp == 0)
        {
            yield return enum1.Current;
            has1 = enum1.MoveNext();
            has2 = enum2.MoveNext();
        }
        else if (cmp < 0) has1 = enum1.MoveNext();
        else has2 = enum2.MoveNext();
    }
}

public IEnumerable<KeyValuePair<K, V>> Except(BaseOrderedMap<K, V, TStrategy> other)
{
    using var enum1 = this.GetEnumerator();
    using var enum2 = other.GetEnumerator();
    
    bool has1 = enum1.MoveNext();
    bool has2 = enum2.MoveNext();

    while (has1 && has2)
    {
        int cmp = _strategy.Compare(enum1.Current.Key, enum2.Current.Key);
        if (cmp == 0)
        {
            has1 = enum1.MoveNext();
            has2 = enum2.MoveNext();
        }
        else if (cmp < 0)
        {
            yield return enum1.Current;
            has1 = enum1.MoveNext();
        }
        else
        {
            has2 = enum2.MoveNext();
        }
    }

    while (has1)
    {
        yield return enum1.Current;
        has1 = enum1.MoveNext();
    }
}

public IEnumerable<KeyValuePair<K, V>> SymmetricExcept(BaseOrderedMap<K, V, TStrategy> other)
{
    using var enum1 = this.GetEnumerator();
    using var enum2 = other.GetEnumerator();
    
    bool has1 = enum1.MoveNext();
    bool has2 = enum2.MoveNext();

    while (has1 && has2)
    {
        int cmp = _strategy.Compare(enum1.Current.Key, enum2.Current.Key);
        if (cmp == 0)
        {
            has1 = enum1.MoveNext();
            has2 = enum2.MoveNext();
        }
        else if (cmp < 0)
        {
            yield return enum1.Current;
            has1 = enum1.MoveNext();
        }
        else
        {
            yield return enum2.Current;
            has2 = enum2.MoveNext();
        }
    }

    while (has1)
    {
        yield return enum1.Current;
        has1 = enum1.MoveNext();
    }
    while (has2)
    {
        yield return enum2.Current;
        has2 = enum2.MoveNext();
    }
}
}