package net.ranides.assira.collection.sets;

import net.ranides.assira.generic.CompareUtils;

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

public class WBTreeSet<K> extends ASortedSet<K> {

    private final WBTreeSetBase<K> tree;

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

    public WBTreeSet(Comparator<? super K> comparator) {
        super(comparator);
        this.tree = new WBTreeSetBase<>(comparator, true);
    }

    @Override
    public boolean contains(Object o) {
        return tree.find((K)o) != null;
    }

    @Override
    public boolean add(K k) {
        int prev = tree.size();
        tree.add(k);
        return tree.size() != prev;
    }

    @Override
    public boolean remove(Object value) {
        return tree.remove((K)value) != null;
    }

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

    @Override
    public ListIterator<K> iterator() {
        return new TreeIterator();
    }

    @Override
    public ListIterator<K> iterator(K begin) {
        WBTreeSetBase.Node<K> after = tree.find(begin, WBTreeSetBase.FindMode.GT);
        int index = WBTreeSetBase.indexOf(after);
        return new TreeIterator(index < 0 ? tree.size() : index);
    }

    @Override
    protected ASortedSet<K> subset(K begin, boolean bottom, K end, boolean top) {
        return new TreeSubset(begin, bottom, end, top);
    }

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

    @Override
    public K first() {
        WBTreeSetBase.Node<K> node = tree.first();
        if(node == null) {
            throw new NoSuchElementException();
        }
        return node.value;
    }

    @Override
    public K last() {
        WBTreeSetBase.Node<K> node = tree.last();
        if(node == null) {
            throw new NoSuchElementException();
        }
        return node.value;
    }

    private class TreeSubset extends ASubset {

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

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

        @Override
        public ListIterator<K> iterator() {
            return new SIterator();
        }

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

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

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

        @Override
        public K first() {
            WBTreeSetBase.Node<K> out = $first();
            if(out == null) {
                throw new NoSuchElementException();
            }
            return out.value;
        }

        @Override
        public K last() {
            WBTreeSetBase.Node<K> out = $last();
            if(out == null) {
                throw new NoSuchElementException();
            }
            return out.value;
        }

        private class SIterator implements ListIterator<K> {

            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;

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

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

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

            @Override
            public final K next() {
                if (!hasNext()) {
                    throw new NoSuchElementException();
                }
                return tree.at((last = index++) + minIndex).value;
            }

            @Override
            public final K previous() {
                if (!hasPrevious()) {
                    throw new NoSuchElementException();
                }
                return tree.at((last = --index) + minIndex).value;
            }

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

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

            @Override
            public final void add(K value) {
                throw new UnsupportedOperationException();
            }

            @Override
            public final void set(K 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<K> {

        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 K next() {
            if (!hasNext()) {
                throw new NoSuchElementException();
            }
            return tree.at(last = index++).value;
        }

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

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

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

        @Override
        public final void add(K value) {
            throw new UnsupportedOperationException();
        }

        @Override
        public final void set(K 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;
        }
    }

}
