package net.ranides.assira.collection.maps;

import net.ranides.assira.generic.CompareUtils;

import java.io.Serializable;
import java.util.Comparator;
import java.util.ListIterator;
import java.util.NoSuchElementException;

public class WBTreeMap<K,V> extends ASortedMap<K,V> {

    private final WBTreeMapBase<K,V> tree;

    public WBTreeMap() {
        this(CompareUtils.comparator());
    }

    public WBTreeMap(Comparator<? super K> comparator) {
        super(comparator);
        this.tree = new WBTreeMapBase<>((Serializable & Comparator<? super K>)(a,b) -> {
            try {
                return comparator.compare(a, b);
            } catch (ClassCastException e) {
                return -1;
            }
        }, true);
    }

    @Override
    protected Entry<K, V> findEntry(K key) {
        return tree.find(key);
    }

    @Override
    protected Entry<K, V> firstEntry() {
        return tree.first();
    }

    @Override
    protected Entry<K, V> lastEntry() {
        return tree.last();
    }

    @Override
    protected ListIterator<Entry<K, V>> entryIterator() {
        return new TreeIterator();
    }

    @Override
    protected ListIterator<Entry<K, V>> entryIterator(K begin) {
        WBTreeMapBase.Node<K, V> after = tree.find(begin, WBTreeMapBase.FindMode.GE);
        int index = WBTreeMapBase.indexOf(after);
        return new TreeIterator(index < 0 ? tree.size() : index);
    }

    @Override
    protected ASortedMap<K, V> submap(K begin, boolean bottom, K end, boolean top) {
        return new TreeSubmap(begin, bottom, end, top);
    }

    @Override
    public int size() {
        return tree.size();
    }

    @Override
    public boolean remove(Object key, Object value) {
        WBTreeMapBase.Node<K, V> prev = tree.find((K)key);
        if(prev != null && CompareUtils.equals(prev.value, value)) {
            tree.removeNode(prev);
            return true;
        }
        return false;
    }

    @Override
    public V put(K key, V value) {
        WBTreeMapBase.Node<K, V> prev = tree.find(key);
        if(prev != null) {
            return prev.setValue(value);
        } else {
            tree.add(key, value);
            return null;
        }
    }

    @Override
    public V remove(Object key) {
        WBTreeMapBase.Node<K, V> prev = tree.remove((K) key);
        return (prev != null) ? prev.value : null;
    }

    @Override
    public void clear() {
        tree.clear();
    }

    private class TreeSubmap extends ASortedMap<K,V>.ASubmap {

        public TreeSubmap(K begin, boolean bottom, K end, boolean top) {
            super(begin, bottom, end, top);
        }

        @Override
        public int size() {
            WBTreeMapBase.Node<K,V> first = $first();
            WBTreeMapBase.Node<K,V> last = $last();
            if(first == null || last == null) {
                return 0;
            } else {
                return WBTreeMapBase.indexOf(last) - WBTreeMapBase.indexOf(first) + 1;
            }
        }

        private WBTreeMapBase.Node<K,V> $first() {
            WBTreeMapBase.Node<K,V> out = bottom ? tree.first() : tree.find(begin, WBTreeMapBase.FindMode.GE);
            if(out == null || !inside(out.key)) {
                return null;
            }
            return out;
        }

        public WBTreeMapBase.Node<K,V> $last() {
            WBTreeMapBase.Node<K,V> out = top ? tree.last() : tree.find(end, WBTreeMapBase.FindMode.LT);
            if(out == null || !inside(out.key)) {
                return null;
            }
            return out;
        }

        @Override
        protected ListIterator<Entry<K,V>> entryIterator() {
            return new SIterator();
        }

        @Override
        protected ListIterator<Entry<K,V>> entryIterator(K from) {
            WBTreeMapBase.Node<K,V> fromNode;
            if(bottom) {
                fromNode = tree.find(from, WBTreeMapBase.FindMode.GT);
            }
            else if(compare(from, begin) > 0) {
                fromNode = tree.find(from, WBTreeMapBase.FindMode.GT);
            } else {
                fromNode = tree.find(begin, WBTreeMapBase.FindMode.GE);
            }
            if(fromNode == null) {
                return new SIterator(size());
            } else {
                return new SIterator(WBTreeMapBase.indexOf(fromNode) - WBTreeMapBase.indexOf($first()));
            }
        }

        @Override
        protected Entry<K,V> findEntry(K key) {
            if (!inside(key) ) {
                return null;
            }
            return WBTreeMap.this.findEntry(key);
        }

        @Override
        protected Entry<K,V> firstEntry() {
            WBTreeMapBase.Node<K,V> out = $first();
            if(out == null) {
                throw new NoSuchElementException();
            }
            return out;
        }

        @Override
        protected Entry<K,V> lastEntry() {
            WBTreeMapBase.Node<K,V> out = $last();
            if(out == null) {
                throw new NoSuchElementException();
            }
            return out;
        }



        private class SIterator implements ListIterator<Entry<K,V>> {

            private int index;
            private int last;

            private final int minIndex;
            private int maxIndex;

            private final int shift;

            public SIterator() {
                this(0);
            }

            public SIterator(int index) {
                this.index = index;
                this.last = -1;
                this.shift = index;

                WBTreeMapBase.Node<K,V> first = $first();
                WBTreeMapBase.Node<K,V> last = $last();
                if(first == null || last == null) {
                    this.minIndex = 0;
                    this.maxIndex = 0;
                } else {
                    this.minIndex = WBTreeMapBase.indexOf(first);
                    this.maxIndex = WBTreeMapBase.indexOf(last) + 1;
                }
            }

            @Override
            public final boolean hasNext() {
                return minIndex+index < maxIndex;
            }

            @Override
            public final boolean hasPrevious() {
                return minIndex+index > minIndex;
            }

            @Override
            public Entry<K,V> next() {
                if (!hasNext()) {
                    throw new NoSuchElementException();
                }
                return tree.at((last = index++) + minIndex);
            }

            @Override
            public Entry<K,V> previous() {
                if (!hasPrevious()) {
                    throw new NoSuchElementException();
                }
                return tree.at((last = --index) + minIndex);
            }

            @Override
            public final int nextIndex() {
                return index - minIndex - shift;
            }

            @Override
            public final int previousIndex() {
                return index - minIndex - shift - 1;
            }

            @Override
            public final void add(Entry<K,V> value) {
                throw new UnsupportedOperationException();
            }

            @Override
            public final void set(Entry<K,V> value) {
                throw new UnsupportedOperationException();
            }

            @Override
            public final void remove() {
                if (last == -1) {
                    throw new IllegalStateException();
                }
                tree.removeNode(tree.at(last + minIndex));
                // If the last operation was a next(),
                // we are removing an element *before* us,
                // and we must decrease pos correspondingly.
                if (last < index) {
                    index--;
                }
                last = -1;
                maxIndex--;
            }
        }

    }

    private class TreeIterator implements ListIterator<Entry<K,V>> {

        private int index;

        private int last;

        private final int offset;

        public TreeIterator() {
            this(0);
        }

        public TreeIterator(int index) {
            this.last = -1;
            this.index = index;
            this.offset = index;
        }

        @Override
        public final boolean hasNext() {
            return index < size();
        }

        @Override
        public final boolean hasPrevious() {
            return index > 0;
        }

        @Override
        public final Entry<K,V> next() {
            if (!hasNext()) {
                throw new NoSuchElementException();
            }
            return tree.at(last = index++);
        }

        @Override
        public final Entry<K,V> previous() {
            if (!hasPrevious()) {
                throw new NoSuchElementException();
            }
            return tree.at(last = --index);
        }

        @Override
        public final int nextIndex() {
            return index - offset;
        }

        @Override
        public final int previousIndex() {
            return index - offset - 1;
        }

        @Override
        public final void add(Entry<K,V> value) {
            throw new UnsupportedOperationException();
        }

        @Override
        public final void set(Entry<K,V> value) {
            throw new UnsupportedOperationException();
        }

        @Override
        public final void remove() {
            if (last == -1) {
                throw new IllegalStateException();
            }
            tree.removeNode(tree.at(last));
            // If the last operation was a next(),
            // we are removing an element *before* us,
            // and we must decrease pos correspondingly.
            if (last < index) {
                index--;
            }
            last = -1;
        }
    }

}
