// Copyright 2026 International Digital Economy Academy
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

/// Tree-local helpers used by vector push and tail normalization.

///|
fn append_sizes_last(sizes : FixedArray[Int]?, delta : Int) -> FixedArray[Int]? {
  match sizes {
    Some(sizes) => {
      let new_sizes = sizes.copy()
      new_sizes[new_sizes.length() - 1] += delta
      Some(new_sizes)
    }
    None => None
  }
}

///|
fn push_sizes_last(sizes : FixedArray[Int]?, delta : Int) -> FixedArray[Int]? {
  match sizes {
    Some(sizes) =>
      Some(immutable_push(sizes, delta + sizes[sizes.length() - 1]))
    None => None
  }
}

///|
/// Sizes for a previously radix (`None`) node whose right spine just absorbed a
/// leaf of `delta` elements. It stays radix only when a full leaf was appended;
/// any partial leaf leaves the node non-full, so it needs explicit sizes.
fn[T] radix_append_sizes(
  children : FixedArray[Tree[T]],
  child_shift : Int,
  delta : Int,
) -> FixedArray[Int]? {
  if delta == BRANCHING_FACTOR {
    None
  } else {
    compute_sizes(children, child_shift)
  }
}

///|
/// Append a leaf to the right spine of a tree without going through generic tree concat.
fn[T] Tree::append_right_leaf(
  self : Tree[T],
  shift : Int,
  leaf : FixedArray[T],
) -> (Tree[T], Int) {
  let delta = leaf.length()
  guard delta > 0 else { return (self, shift) }

  fn worker(
    node : Tree[T],
    shift : Int,
    leaf : FixedArray[T],
    delta : Int,
  ) -> Tree[T]? {
    match node {
      Leaf(elems) => {
        guard shift == 0 else {
          abort("Unreachable: Leaf should only appear at shift 0")
        }
        if elems.length() + delta <= BRANCHING_FACTOR {
          Some(Leaf(immutable_concat(elems, leaf)))
        } else {
          None
        }
      }
      Node(nodes, sizes) => {
        let len = nodes.length()
        match worker(nodes[len - 1], shift - NUM_BITS, leaf, delta) {
          Some(new_node) => {
            // The right spine absorbed the leaf below us. If `sizes` was `None`
            // (radix) and we just folded in a partial leaf, the node is no
            // longer full and must carry an explicit sizes array.
            let new_nodes = immutable_set(nodes, len - 1, new_node)
            let new_sizes = match sizes {
              Some(_) => append_sizes_last(sizes, delta)
              None => radix_append_sizes(new_nodes, shift - NUM_BITS, delta)
            }
            Some(Node(new_nodes, new_sizes))
          }
          None =>
            if len < BRANCHING_FACTOR {
              let new_nodes = immutable_push(
                nodes,
                new_branch_left(leaf, shift - NUM_BITS),
              )
              let new_sizes = match sizes {
                Some(_) => push_sizes_last(sizes, delta)
                None => radix_append_sizes(new_nodes, shift - NUM_BITS, delta)
              }
              Some(Node(new_nodes, new_sizes))
            } else {
              None
            }
        }
      }
      Empty => Some(Leaf(leaf))
    }
  }

  match worker(self, shift, leaf, delta) {
    Some(new_tree) => (new_tree, shift)
    None => {
      // A fresh root over `[self, new_branch]` stays radix only when every leaf
      // is full: `self` is already radix and the appended leaf is full too.
      let new_branch = new_branch_left(leaf, shift)
      let self_radix = match self {
        Node(_, None) => true
        Leaf(elems) => elems.length() == BRANCHING_FACTOR
        Node(_, Some(_)) => false
        Empty =>
          abort("Unreachable: Empty tree should have been handled in worker")
      }
      let new_sizes : FixedArray[Int]? = if self_radix &&
        delta == BRANCHING_FACTOR {
        None
      } else {
        let len = self.size(shift)
        Some([len, len + delta])
      }
      (Node([self, new_branch], new_sizes), shift + NUM_BITS)
    }
  }
}

///|
/// Concatenate two trees while threading an extra right-most leaf from the left side
/// through the splice recursion, avoiding a full normalize-then-concat pass.
fn[T] Tree::concat_with_suffix(
  left : Tree[T],
  left_shift : Int,
  suffix : FixedArray[T],
  right : Tree[T],
  right_shift : Int,
  top : Bool,
) -> (Tree[T], Int) {
  if suffix.is_empty() {
    return Tree::concat(left, left_shift, right, right_shift, top)
  }
  if left_shift > right_shift {
    let (c, c_shift) = Tree::concat_with_suffix(
      left.right_child(),
      left_shift - NUM_BITS,
      suffix,
      right,
      right_shift,
      false,
    )
    guard! c_shift == left_shift
    rebalance(left, c, Empty, left_shift, top)
  } else if right_shift > left_shift {
    let (c, c_shift) = Tree::concat_with_suffix(
      left,
      left_shift,
      suffix,
      right.left_child(),
      right_shift - NUM_BITS,
      false,
    )
    guard! c_shift == right_shift
    rebalance(Empty, c, right, right_shift, top)
  } else if left_shift == 0 {
    let leaves = FixedArray::from_array([left, Leaf(suffix), right])
    let (node_counts, new_len) = redis_plan(leaves)
    let new_children = redis(leaves, node_counts, new_len, 0)
    if top && new_len == 1 {
      (new_children[0], 0)
    } else {
      let node = Node(new_children, compute_sizes(new_children, 0))
      (node, NUM_BITS)
    }
  } else {
    let (c, c_shift) = Tree::concat_with_suffix(
      left.right_child(),
      left_shift - NUM_BITS,
      suffix,
      right.left_child(),
      right_shift - NUM_BITS,
      false,
    )
    guard! c_shift == left_shift
    guard! c_shift == right_shift
    rebalance(left, c, right, left_shift, top)
  }
}