package avl

import (
	"github.com/emirpasic/gods/stacks/arraystack"
)

type Iterator struct {
	op *AVL

	dir    int
	up     *Node
	cur    *Node
	tstack *arraystack.Stack
	// curnext *Node
}

func initIterator(avltree *AVL) *Iterator {
	iter := &Iterator{op: avltree, tstack: arraystack.New()}
	iter.up = avltree.root
	return iter
}

func (iter *Iterator) Value() interface{} {
	return iter.cur.value
}

func (iter *Iterator) Next() (result bool) {

	if iter.dir < 1 {
		if iter.dir == -1 {
			if iter.cur != nil {
				iter.tstack.Clear()
				iter.curPushNextStack(iter.cur)

				iter.up = iter.cur
				for iter.up != nil {
					if iter.up.child == 1 {
						iter.up = iter.up.parent
						break
					} else {
						iter.up = iter.up.parent
					}
				}

			}
		}
		iter.dir = 1
	}

	if iter.tstack.Size() == 0 {
		if iter.up != nil {
			iter.tstack.Push(iter.up)
			for iter.up != nil {
				if iter.up.child == 1 {
					iter.up = iter.up.parent
					break
				} else {
					iter.up = iter.up.parent
				}
			}
		} else {
			return false
		}
	}

	if v, ok := iter.tstack.Pop(); ok {
		iter.cur = v.(*Node)
		iter.curPushNextStack(iter.cur)
		return true
	}

	return false
}

func (iter *Iterator) curPushNextStack(cur *Node) {
	next := cur.children[0]

	if next != nil {
		iter.tstack.Push(next)
		for next.children[1] != nil {
			next = next.children[1]
			iter.tstack.Push(next)
		}
	}
}

func (iter *Iterator) curPushPrevStack(cur *Node) {
	prev := cur.children[1]

	if prev != nil {
		iter.tstack.Push(prev)
		for prev.children[0] != nil {
			prev = prev.children[0]
			iter.tstack.Push(prev)
		}
	}
}

func (iter *Iterator) Prev() (result bool) {

	if iter.dir > -1 {

		if iter.dir == 1 {
			if iter.cur != nil {
				iter.tstack.Clear()
				iter.curPushPrevStack(iter.cur)

				iter.up = iter.cur
				for iter.up != nil {
					if iter.up.child == 0 {
						iter.up = iter.up.parent
						break
					} else {
						iter.up = iter.up.parent
					}
				}
			}
		}
		iter.dir = -1
	}

	if iter.tstack.Size() == 0 {
		if iter.up != nil {
			iter.tstack.Push(iter.up)
			for iter.up != nil {
				if iter.up.child == 0 {
					iter.up = iter.up.parent
					break
				} else {
					iter.up = iter.up.parent
				}
			}
		} else {
			return false
		}
	}

	if v, ok := iter.tstack.Pop(); ok {
		iter.cur = v.(*Node)
		iter.curPushPrevStack(iter.cur)
		return true
	}

	return false
}