import scala.collection.mutable

class AVLTree[A](implicit val ordering: Ordering[A]) extends mutable.SortedSet[A] {

  if (ordering eq null) throw new NullPointerException("ordering must not be null")

  private var _root: AVLNode = _
  private var _size = 0

  override def size: Int = _size

  override def foreach[U](f: A => U): Unit = {
    val stack = mutable.Stack[AVLNode]()
    var current = root
    var done = false

    while (!done) {
      if (current != null) {
        stack.push(current)
        current = current.left
      } else if (stack.nonEmpty) {
        current = stack.pop()
        f.apply(current.key)

        current = current.right
      } else {
        done = true
      }
    }
  }

  def root: AVLNode = _root

  override def isEmpty: Boolean = root == null

  override def min[B >: A](implicit cmp: Ordering[B]): A = minNode().key

  def minNode(): AVLNode = {
    if (root == null) throw new UnsupportedOperationException("empty tree")
    var node = root
    while (node.left != null) node = node.left
    node
  }

  override def max[B >: A](implicit cmp: Ordering[B]): A = maxNode().key

  def maxNode(): AVLNode = {
    if (root == null) throw new UnsupportedOperationException("empty tree")
    var node = root
    while (node.right != null) node = node.right
    node
  }

  def next(node: AVLNode): Option[AVLNode] = {
    var successor = node
    if (successor != null) {
      if (successor.right != null) {
        successor = successor.right
        while (successor != null && successor.left != null) {
          successor = successor.left
        }
      } else {
        successor = node.parent
        var n = node
        while (successor != null && successor.right == n) {
          n = successor
          successor = successor.parent
        }
      }
    }
    Option(successor)
  }

  def prev(node: AVLNode): Option[AVLNode] = {
    var predecessor = node
    if (predecessor != null) {
      if (predecessor.left != null) {
        predecessor = predecessor.left
        while (predecessor != null && predecessor.right != null) {
          predecessor = predecessor.right
        }
      } else {
        predecessor = node.parent
        var n = node
        while (predecessor != null && predecessor.left == n) {
          n = predecessor
          predecessor = predecessor.parent
        }
      }
    }
    Option(predecessor)
  }

  override def rangeImpl(from: Option[A], until: Option[A]): mutable.SortedSet[A] = ???

  override def +=(key: A): AVLTree.this.type = {
    insert(key)
    this
  }

  def insert(key: A): AVLNode = {
    if (root == null) {
      _root = new AVLNode(key)
      _size += 1
      return root
    }

    var node = root
    var parent: AVLNode = null
    var cmp = 0

    while (node != null) {
      parent = node
      cmp = ordering.compare(key, node.key)
      if (cmp == 0) return node