class_name BST

class _Node extends RefCounted:
    var value: int
    var left: Option = Option.none
    var right: Option = Option.none

    func _init(value: int) -> void:
        self.value = value

    func size() -> int:
        return BST.size(self)

    func has_children() -> bool:
        return BST.has_children(self)

    func balance() -> _Node:
        return BST.balance(self).unwrap()

    func insert(value: int) -> _Node:
        return BST.insert(self, value)

    func search(value: int) -> Option:
        return BST.search(self, value)

    func remove(value: int) -> Option:
        return BST.remove(self, value)

    func iter() -> Iterator:
        return BST.iter(self)

static func from(value: int) -> _Node:
    return _Node.new(value)

static func size(node: _Node) -> int:
    var left = node.left.map(size).or_default(0)
    var right = node.right.map(size).or_default(0)
    return left + right + 1

static func _store_in_order(_root: Option, nodes: Array[_Node] = []) -> Array[_Node]:
    if _root.is_none(): return nodes

    var root = _root.unwrap()
    _store_in_order(root.left, nodes)
    nodes.append(root)
    _store_in_order(root.right, nodes)
    return nodes

static func _build_balanced_tree(nodes: Array[_Node], start: int = 0, end: int = len(nodes) - 1) -> Option:
    if start > end:
        return Option.none

    var mid: int = (start + end) / 2
    var root = nodes[mid]

    root.left = _build_balanced_tree(nodes, start, mid - 1)
    root.right = _build_balanced_tree(nodes, mid + 1, end)

    return Option.some(root)

static func balance(root: _Node) -> Option:
    return _build_balanced_tree(_store_in_order(Option.some(root)))

static func has_children(node: _Node) -> bool:
    return node.left.is_none() and node.right.is_none()

static func _insert(_root: Option, node: _Node) -> _Node:
    if _root.is_none(): return node

    var root = _root.unwrap()
    if root.value == node.value: return root
    elif root.value < node.value:
        root.right = Option.some(_insert(root.right, node))
    else:
        root.left = Option.some(_insert(root.left, node))
    return root

static func insert(root: _Node, value: int) -> _Node:
    return _insert(Option.some(root), _Node.new(value))

static func _search(_root: Option, value: int) -> Option:
    if _root.is_none():
        return Option.none
    var root = _root.unwrap()
    if root.value == value:
        return _root
    if root.value < value:
        return _search(root.right, value)
    return _search(root.left, value)

static func search(root: _Node, value: int) -> Option:
    return _search(Option.some(root), value)

static func _get_successor(node: _Node) -> _Node:
    var current = node.right
    while current.is_some() and current.left.is_some():
        current = node.left.unwrap()
    return current

static func _remove(_root: Option, value: int) -> Option:
    if _root.is_none(): return _root

    var root = _root.unwrap()
    if root.value < value:
        return _remove(root.right, value)
    elif root.value > value:
        return _remove(root.left, value)
    else:
        if root.right.is_none():
            return root.left
        elif root.left.is_none():
            return root.right

        var successor = _get_successor(root)
        root.value = successor.value
        root.right = remove(root.right, successor.value)

    return Option.some(root)

static func remove(root: _Node, value: int) -> Option:
    return _remove(Option.some(root), value)

static func traverse(root: _Node, fn: Callable):
    root.left.map(fn.call)
    fn.call(root.value)
    root.right.map(fn.call)

static func iter(node: _Node) -> Iterator:
    var iters = Option.collect_some([
        node.left.map(iter),
        Option.some(SingletonIterator.new(node.value)),
        node.right.map(iter)
    ], true).unwrap()

    var iter_arr: Array[Iterator]
    iter_arr.assign(iters)

    return FusedIterator.new(iter_arr)