module datatypes

// Internal representation of the tree node
@[heap]
struct BSTreeNode[T] {
mut:
    // Mark a node as initialized
    is_init bool
    // Value of the node
    value T
    // The parent of the node
    parent &BSTreeNode[T] = unsafe { 0 }
    // The left side with value less than the
    // value of this node
    left &BSTreeNode[T] = unsafe { 0 }
    // The right side with value grater than the
    // value of thiss node
    right &BSTreeNode[T] = unsafe { 0 }
}

// Create new root bst node
fn new_root_node[T](value T) &BSTreeNode[T] {
    return &BSTreeNode[T]{
        is_init: true
        value:   value
        parent:  new_none_node[T](true)
        left:    new_none_node[T](false)
        right:   new_none_node[T](false)
    }
}

// new_node creates a new bst node with a parent reference.
fn new_node[T](parent &BSTreeNode[T], value T) &BSTreeNode[T] {
    return &BSTreeNode[T]{
        is_init: true
        value:   value
        parent:  parent
    }
}

// new_none_node creates a dummy node.
fn new_none_node[T](init bool) &BSTreeNode[T] {
    return &BSTreeNode[T]{
        is_init: init
    }
}

// bind to an actual instance of a node.
fn (mut node BSTreeNode[T]) bind(mut to_bind BSTreeNode[T], _left bool) {
    node.left = to_bind.left
    node.right = to_bind.right
    node.value = to_bind.value
    node.is_init = to_bind.is_init
    to_bind = *new_none_node[T](false)
}

// Pure Binary Seach Tree implementation
//
// Pure V implementation of the Binary Search Tree
// Time complexity of main operation O(log N)
// Space complexity O(N)
pub struct BSTree[T] {
mut:
    root &BSTreeNode[T] = unsafe { 0 }
}

// insert give the possibility to insert an element in the BST.
pub fn (mut bst BSTree[T]) insert(value T) bool {
    if bst.is_empty() {
        bst.root = new_root_node(value)
        return true
    }
    return bst.insert_helper(mut bst.root, value)
}

// insert_helper walks the tree and inserts the given node.
fn (mut bst BSTree[T]) insert_helper(mut node BSTreeNode[T], value T) bool {
    if node.value < value {
        if unsafe { node.right != 0 } && node.right.is_init {
            return bst.insert_helper(mut node.right, value)
        }
        node.right = new_node(node, value)
        return true
    } else if node.value > value {
        if unsafe { node.left != 0 } && node.left.is_init {
            return bst.insert_helper(mut node.left, value)
        }
        node.left = new_node(node, value)
        return true
    }
    return false
}

// contains checks if an element with a given `value` is inside the BST.
pub fn (bst &BSTree[T]) contains(value T) bool {
    return bst.contains_helper(bst.root, value)
}

// contains_helper is a helper function to walk the tree, and return
// the absence or presence of the `value`.
fn (bst &BSTree[T]) contains_helper(node &BSTreeNode[T], value T) bool {
    if unsafe { node == 0 } || !node.is_init {
        return false
    }
    if node.value < value {
        return bst.contains_helper(node.right, value)
    } else if node.value > value {
        return bst.contains_helper(node.left, value)
    }
    assert node.value == value
    return true
}

// remove removes an element with `value` from the BST.
pub fn (mut bst BSTree[T]) remove(value T) bool {
    if bst.is_empty() {
        return false
    }
    return bst.remove_helper(mut bst.root, value, false)
}

fn (mut bst BSTree[T]) remove_helper(mut node BSTreeNode[T], value T, left bool) bool {
    if !node.is_init {
        return false
    }
    if node.value == value {
        if unsafe { node.left != 0 } && node.left.is_init {
            // In order to remove the element we need to bring up as parent the max of the
            // left sub-tree.
            mut max_node := bst.get_max_from_right(node.left)
            node.bind(mut max_node, true)
        } else if unsafe { node.right != 0 } && node.right.is_init {
            // Bring up the element with the minimum value in the right sub-tree.
            mut min_node := bst.get_min_from_left(node.right)
            node.bind(mut min_node, false)
        } else {
            mut parent := node.parent
            if left {
                parent.left = new_none_node[T](false)
            } else {
                parent.right = new_none_node[T](false)
            }
            node = *new_none_node[T](false)
        }
        return true
    }

    if node.value < value {
        return bst.remove_helper(mut node.right, value, false)
    }
    return bst.remove_helper(mut node.left, value, true)
}

// get_max_from_right returns the max element of the BST following the right branch.
fn (bst &BSTree[T]) get_max_from_right(node &BSTreeNode[T]) &BSTreeNode[T] {
    if unsafe { node == 0 } {
        return new_none_node[T](false)
    }
    right_node := node.right
    if unsafe { right_node == 0 } || !right_node.is_init {
        return node
    }
    return bst.get_max_from_right(right_node)
}

// get_min_from_left returns the min element of the BST by following the left branch.
fn (bst &BSTree[T]) get_min_from_left(node &BSTreeNode[T]) &BSTreeNode[T] {
    if unsafe { node == 0 } {
        return new_none_node[T](false)
    }
    left_node := node.left
    if unsafe { left_node == 0 } || !left_node.is_init {
        return node
    }
    return bst.get_min_from_left(left_node)
}

// is_empty checks if the BST is empty
pub fn (bst &BSTree[T]) is_empty() bool {
    return unsafe { bst.root == 0 }
}

// in_order_traversal traverses the BST in order, and returns the result as an array.
pub fn (bst &BSTree[T]) in_order_traversal() []T {
    mut result := []T{}
    bst.in_order_traversal_helper(bst.root, mut result)
    return result
}

// in_order_traversal_helper helps traverse the BST, and accumulates the result in the `result` array.
fn (bst &BSTree[T]) in_order_traversal_helper(node &BSTreeNode[T], mut result []T) {
    if unsafe { node == 0 } || !node.is_init {
        return
    }
    bst.in_order_traversal_helper(node.left, mut result)
    result << node.value
    bst.in_order_traversal_helper(node.right, mut result)
}

// post_order_traversal traverses the BST in post order, and returns the result in an array.
pub fn (bst &BSTree[T]) post_order_traversal() []T {
    mut result := []T{}
    bst.post_order_traversal_helper(bst.root, mut result)
    return result
}

// post_order_traversal_helper is a helper function that traverses the BST in post order,
// accumulating the result in an array.
fn (bst &BSTree[T]) post_order_traversal_helper(node &BSTreeNode[T], mut result []T) {
    if unsafe { node == 0 } || !node.is_init {
        return
    }

    bst.post_order_traversal_helper(node.left, mut result)
    bst.post_order_traversal_helper(node.right, mut result)
    result << node.value
}

// pre_order_traversal traverses the BST in pre order, and returns the result as an array.
pub fn (bst &BSTree[T]) pre_order_traversal() []T {
    mut result := []T{}
    bst.pre_order_traversal_helper(bst.root, mut result)
    return result
}

// pre_order_traversal_helper is a helper function to traverse the BST
// in pre order and accumulates the results in an array.
fn (bst &BSTree[T]) pre_order_traversal_helper(node &BSTreeNode[T], mut result []T) {
    if unsafe { node == 0 } || !node.is_init {
        return
    }
    result << node.value
    bst.pre_order_traversal_helper(node.left, mut result)
    bst.pre_order_traversal_helper(node.right, mut result)
}

// get_node is a helper method to ge the internal representation of the node with the `value`.
fn (bst &BSTree[T]) get_node(node &BSTreeNode[T], value T) &BSTreeNode[T] {
    if unsafe { node == 0 } || !node.is_init {
        return new_none_node[T](false)
    }
    if node.value == value {
        return node
    }

    if node.value < value {
        return bst.get_node(node.right, value)
    }
    return bst.get_node(node.left, value)
}

// to_left returns the value of the node to the left of the node with `value` specified if it exists
pub fn (bst &BSTree[T]) to_left(value T) !T {
    if bst.is_empty() {
        return error('BSTree is empty')
    }
    node := bst.get_node(bst.root, value)
    if !node.is_init {
        return error('BSTree is not initialized')
    }
    left_node := node.left
    if unsafe { left_node == 0 } || !left_node.is_init {
        return error('No left node exists for the given value')
    }
    return left_node.value
}

// to_right returns the value of the node to the right of the node with `value` specified if it exists
pub fn (bst &BSTree[T]) to_right(value T) !T {
    if bst.is_empty() {
        return error('BSTree is empty')
    }
    node := bst.get_node(bst.root, value)
    if !node.is_init {
        return error('BSTree is not initialized')
    }
    right_node := node.right
    if unsafe { right_node == 0 } || !right_node.is_init {
        return error('No right node exists for the given value')
    }
    return right_node.value
}

// max return the max element inside the BST.
// Time complexity O(N) if the BST is not balanced
pub fn (bst &BSTree[T]) max() !T {
    if bst.is_empty() {
        return error('BSTree is empty')
    }
    max := bst.get_max_from_right(bst.root)
    if !max.is_init {
        return error('BSTree is not initialized')
    }
    return max.value
}

// min return the minimum element in the BST.
// Time complexity O(N) if the BST is not balanced.
pub fn (bst &BSTree[T]) min() !T {
    if bst.is_empty() {
        return error('BSTree is empty')
    }
    min := bst.get_min_from_left(bst.root)
    if !min.is_init {
        return error('BSTree is not initialized')
    }
    return min.value
}