///| Tree-batch planning turns a `DraftTree` into explicit target-model query

///| contexts. An inference adapter may concatenate these contexts into one

///| batched forward pass, while this package keeps the planning independent of

///|
/// tensor runtimes and model-specific cache formats.
pub enum TreeBatchError {
  InvalidTree
  MissingNode(Int)
  DuplicateScore(Int)
  MissingScore(Int)
  InvalidScore(Int)
} derive(Eq, Debug)

///|
/// Context excludes the candidate token identified by node_id. Returned
/// logits predict that token, not its successor.
pub struct TreeQuery {
  node_id : Int
  parent_id : Int?
  depth : Int
  context : Array[Int]
}

///|
pub fn TreeQuery::node_id(self : TreeQuery) -> Int {
  self.node_id
}

///|
pub fn TreeQuery::context(self : TreeQuery) -> Array[Int] {
  self.context
}

///| Tree queries arranged by depth. Layers let a backend choose between a

///|
/// fully flattened batch and cache-aware level-by-level execution.
pub struct TreeBatchPlan {
  layers : Array[Array[TreeQuery]]
  query_count : Int
}

///|
pub fn TreeBatchPlan::query_count(self : TreeBatchPlan) -> Int {
  self.query_count
}

///|
pub fn TreeBatchPlan::layer_count(self : TreeBatchPlan) -> Int {
  self.layers.length()
}

///|
pub fn TreeBatchPlan::layer(
  self : TreeBatchPlan,
  depth : Int,
) -> Array[TreeQuery] {
  if depth <= 0 || depth > self.layers.length() {
    []
  } else {
    self.layers[depth - 1]
  }
}

///|
pub fn TreeBatchPlan::flatten(self : TreeBatchPlan) -> Array[TreeQuery] {
  let queries : Array[TreeQuery] = []
  for layer in self.layers {
    for query in layer {
      queries.push(query)
    }
  }
  queries
}

///| Reconstruct the complete token context for a node. It is separate from

///| `DraftTree::path_to` because a target adapter needs the original prompt as

///|
/// well as the tree path.
pub fn context_for_node(
  tree : DraftTree,
  node_id : Int,
) -> Result[Array[Int], TreeBatchError] {
  let node = match tree.node(node_id) {
    Some(value) => value
    None => return Err(MissingNode(node_id))
  }
  let context : Array[Int] = []
  for token in tree.prefix {
    context.push(token)
  }
  if node.parent is Some(parent) {
    for token in tree.path_to(parent) {
      context.push(token)
    }
  }
  Ok(context)
}

///| Plan parent-before-child target queries. Tree validation prevents missing

///|
/// ancestors and malformed depths from leaking into a model backend.
pub fn plan_tree_batch(
  tree : DraftTree,
) -> Result[TreeBatchPlan, TreeBatchError] {
  match validate_tree(tree) {
    Err(_) => return Err(InvalidTree)
    Ok(_) => ()
  }
  let layers : Array[Array[TreeQuery]] = []
  for depth in 1..<(tree.max_depth() + 1) {
    let layer : Array[TreeQuery] = []
    for node in tree.nodes {
      if node.depth == depth {
        let context = match context_for_node(tree, node.id) {
          Ok(value) => value
          Err(error) => return Err(error)
        }
        layer.push({
          node_id: node.id,
          parent_id: node.parent,
          depth: node.depth,
          context,
        })
      }
    }
    layers.push(layer)
  }
  Ok({ query_count: tree.node_count(), layers })
}

///| Target score material returned by an adapter. A score is associated with a

///| node id instead of a positional array so flattened and layered backends

///|
/// cannot accidentally misalign a sibling branch.
pub struct NodeScore {
  node_id : Int
  distribution : Array[Double]
}

///|
pub fn NodeScore::new(node_id : Int, distribution : Array[Double]) -> NodeScore {
  { node_id, distribution }
}

///| Validate and reorder backend scores to tree node order. Verifiers can then

///| safely index the output alongside tree nodes. This detects duplicated,

///|
/// omitted, and malformed scores at the boundary with model code.
pub fn align_node_scores(
  tree : DraftTree,
  scores : Array[NodeScore],
) -> Result[Array[Array[Double]], TreeBatchError] {
  match validate_tree(tree) {
    Err(_) => return Err(InvalidTree)
    Ok(_) => ()
  }
  let aligned : Array[Array[Double]] = []
  for node in tree.nodes {
    let mut found : Array[Double]? = None
    for score in scores {
      if score.node_id == node.id {
        if found is Some(_) {
          return Err(DuplicateScore(node.id))
        }
        if !valid_node_distribution(score.distribution) ||
          score.distribution.length() != node.distribution.length() {
          return Err(InvalidScore(node.id))
        }
        found = Some(score.distribution)
      }
    }
    match found {
      Some(distribution) => aligned.push(distribution)
      None => return Err(MissingScore(node.id))
    }
  }
  for score in scores {
    if tree.node(score.node_id) is None {
      return Err(MissingNode(score.node_id))
    }
  }
  Ok(aligned)
}

///|
/// Render a stable, model-free execution plan useful in CI fixture diffs.
pub fn TreeBatchPlan::render(self : TreeBatchPlan) -> String {
  let mut text = "tree batch queries=" + self.query_count.to_string() + "\n"
  for depth in 1..<(self.layers.length() + 1) {
    let layer = self.layer(depth)
    text = text + "depth=" + depth.to_string() + " nodes="
    for index in 0.. 0 {
        text = text + ","
      }
      text = text + layer[index].node_id.to_string()
    }
    text = text + "\n"
  }
  text
}