///|
/// Context-sensitive model boundary. A provider returns next-token logits for
/// the supplied prefix. This reference backend makes sequential calls; it
/// does not claim to implement GPU tree attention or a batched forward pass.
pub enum ModelTreeError {
  InvalidConfiguration
  ProviderFailure(String)
  BatchResultLengthMismatch
  InvalidLogits
  InvalidTree
  VocabularyMismatch
} derive(Eq, Debug)

///|
pub struct ModelDraftTree {
  tree : DraftTree
  provider_calls : Int
}

///|
pub struct ScoredTree {
  distributions : Array[Array[Double]]
  provider_calls : Int
}

///|
/// Query each parent prefix separately, including distinct siblings' paths.
/// Width is capped by the vocabulary and the global node budget.
pub fn TreePolicy::build_from_model(
  self : TreePolicy,
  prefix : Array[Int],
  depth : Int,
  model : (Array[Int]) -> Result[Array[Double], String],
) -> Result[ModelDraftTree, ModelTreeError] {
  build_model_tree(prefix, depth, self.max_nodes, _ => Ok(self.width), model)
}

///|
pub fn AdaptiveTreePolicy::build_from_model(
  self : AdaptiveTreePolicy,
  prefix : Array[Int],
  model : (Array[Int]) -> Result[Array[Double], String],
) -> Result[ModelDraftTree, ModelTreeError] {
  build_model_tree(
    prefix,
    self.depth,
    self.max_nodes,
    logits => {
      match self.width_for_logits(logits) {
        Ok(width) => Ok(width)
        Err(_) => Err(InvalidLogits)
      }
    },
    model,
  )
}

///|
/// Build a draft tree with one batch-provider call per tree level. Candidate
/// prefixes for the next level are known only after this level is scored, so a
/// single all-level draft request would be invalid.
pub fn TreePolicy::build_from_batch_model(
  self : TreePolicy,
  prefix : Array[Int],
  depth : Int,
  model : (Array[Array[Int]]) -> Result[Array[Array[Double]], String],
) -> Result[ModelDraftTree, ModelTreeError] {
  build_model_tree_batched(
    prefix,
    depth,
    self.max_nodes,
    _ => Ok(self.width),
    model,
  )
}

///|
/// Entropy-adaptive counterpart of `TreePolicy::build_from_batch_model`.
/// Branching width is computed separately for each parent row in a batch.
pub fn AdaptiveTreePolicy::build_from_batch_model(
  self : AdaptiveTreePolicy,
  prefix : Array[Int],
  model : (Array[Array[Int]]) -> Result[Array[Array[Double]], String],
) -> Result[ModelDraftTree, ModelTreeError] {
  self.build_from_batch_model_with_depth_limit(prefix, self.depth, model)
}

///|
/// Build an adaptive draft tree while capping the current depth. Decoders use
/// this at an output boundary so they never score suffix tokens that cannot be
/// emitted in the requested budget.
pub fn AdaptiveTreePolicy::build_from_batch_model_with_depth_limit(
  self : AdaptiveTreePolicy,
  prefix : Array[Int],
  depth_limit : Int,
  model : (Array[Array[Int]]) -> Result[Array[Array[Double]], String],
) -> Result[ModelDraftTree, ModelTreeError] {
  build_model_tree_batched(
    prefix,
    self.depth.min(depth_limit),
    self.max_nodes,
    logits => {
      match self.width_for_logits(logits) {
        Ok(width) => Ok(width)
        Err(_) => Err(InvalidLogits)
      }
    },
    model,
  )
}

///|
fn build_model_tree(
  prefix : Array[Int],
  depth : Int,
  max_nodes : Int,
  width_for : (Array[Double]) -> Result[Int, ModelTreeError],
  model : (Array[Int]) -> Result[Array[Double], String],
) -> Result[ModelDraftTree, ModelTreeError] {
  if depth <= 0 || max_nodes <= 0 || max_nodes > 1024 {
    return Err(InvalidConfiguration)
  }
  let nodes : Array[TreeNode] = []
  let mut frontier : Array[(Int?, Array[Int])] = [(None, prefix.copy())]
  let mut calls = 0
  let mut vocabulary = 0
  for level in 1..<(depth + 1) {
    let next : Array[(Int?, Array[Int])] = []
    for entry in frontier {
      let (parent, context) = entry
      if nodes.length() >= max_nodes {
        break
      }
      let logits = match model(context.copy()) {
        Ok(value) => value
        Err(error) => return Err(ProviderFailure(error))
      }
      calls = calls + 1
      if vocabulary == 0 {
        vocabulary = logits.length()
      }
      if logits.length() != vocabulary {
        return Err(VocabularyMismatch)
      }
      let width = match width_for(logits) {
        Ok(value) => value
        Err(error) => return Err(error)
      }
      let candidates = match top_tokens(logits, width) {
        Ok(value) => value
        Err(_) => return Err(InvalidLogits)
      }
      for candidate in candidates {
        let (token, distribution) = candidate
        if nodes.length() >= max_nodes {
          break
        }
        let id = nodes.length()
        nodes.push({ id, parent, token, distribution, depth: level })
        let extended = context.copy()
        extended.push(token)
        next.push((Some(id), extended))
      }
    }
    frontier = next
    if frontier.is_empty() || nodes.length() >= max_nodes {
      break
    }
  }
  match make_tree(prefix.copy(), nodes) {
    Ok(tree) => Ok({ tree, provider_calls: calls })
    Err(_) => Err(InvalidTree)
  }
}

///|
/// Batch implementation of context-sensitive draft expansion. It preserves
/// the same parent-before-child node order as the sequential builder while
/// making the available level parallelism observable at the provider boundary.
fn build_model_tree_batched(
  prefix : Array[Int],
  depth : Int,
  max_nodes : Int,
  width_for : (Array[Double]) -> Result[Int, ModelTreeError],
  model : (Array[Array[Int]]) -> Result[Array[Array[Double]], String],
) -> Result[ModelDraftTree, ModelTreeError] {
  if depth <= 0 || max_nodes <= 0 || max_nodes > 1024 {
    return Err(InvalidConfiguration)
  }
  let nodes : Array[TreeNode] = []
  let mut frontier : Array[(Int?, Array[Int])] = [(None, prefix.copy())]
  let mut calls = 0
  let mut vocabulary = 0
  for level in 1..<(depth + 1) {
    if frontier.is_empty() || nodes.length() >= max_nodes {
      break
    }
    let contexts : Array[Array[Int]] = []
    for entry in frontier {
      let (_, context) = entry
      contexts.push(context.copy())
    }
    let logits_rows = match model(contexts) {
      Ok(value) => value
      Err(error) => return Err(ProviderFailure(error))
    }
    if logits_rows.length() != frontier.length() {
      return Err(BatchResultLengthMismatch)
    }
    calls = calls + 1
    let next : Array[(Int?, Array[Int])] = []
    for index in 0..= max_nodes {
        break
      }
      let (parent, context) = frontier[index]
      let logits = logits_rows[index]
      if vocabulary == 0 {
        vocabulary = logits.length()
      }
      if logits.length() != vocabulary {
        return Err(VocabularyMismatch)
      }
      let width = match width_for(logits) {
        Ok(value) => value
        Err(error) => return Err(error)
      }
      let candidates = match top_tokens(logits, width) {
        Ok(value) => value
        Err(_) => return Err(InvalidLogits)
      }
      for candidate in candidates {
        if nodes.length() >= max_nodes {
          break
        }
        let (token, distribution) = candidate
        let id = nodes.length()
        nodes.push({ id, parent, token, distribution, depth: level })
        let extended = context.copy()
        extended.push(token)
        next.push((Some(id), extended))
      }
    }
    frontier = next
  }
  match make_tree(prefix.copy(), nodes) {
    Ok(tree) => Ok({ tree, provider_calls: calls })
    Err(_) => Err(InvalidTree)
  }
}

///|
/// Evaluate each unique parent context once, then share its distribution
/// among siblings. provider_calls counts real callback invocations.
pub fn score_tree_from_model(
  tree : DraftTree,
  model : (Array[Int]) -> Result[Array[Double], String],
) -> Result[ScoredTree, ModelTreeError] {
  match validate_tree(tree) {
    Err(_) => return Err(InvalidTree)
    Ok(_) => ()
  }
  let parents : Array[Int?] = []
  let rows : Array[Array[Double]] = []
  let distributions : Array[Array[Double]] = []
  for node in tree.nodes {
    let mut found : Array[Double]? = None
    for i in 0.. row
      None => {
        let context = match context_for_node(tree, node.id) {
          Ok(value) => value
          Err(_) => return Err(InvalidTree)
        }
        let logits = match model(context) {
          Ok(value) => value
          Err(error) => return Err(ProviderFailure(error))
        }
        if logits.length() != node.distribution.length() {
          return Err(VocabularyMismatch)
        }
        let row = match softmax(logits) {
          Ok(value) => value
          Err(_) => return Err(InvalidLogits)
        }
        parents.push(node.parent)
        rows.push(row)
        row
      }
    }
    distributions.push(row)
  }
  Ok({ distributions, provider_calls: parents.length() })
}

///|
/// Collect one context for every unique parent in node order. A sibling group
/// shares a next-token distribution, so sending duplicate contexts to a model
/// backend only wastes work.
fn unique_parent_queries(
  tree : DraftTree,
  queries : Array[TreeQuery],
) -> Result[(Array[Int?], Array[Array[Int]]), ModelTreeError] {
  let parents : Array[Int?] = []
  let contexts : Array[Array[Int]] = []
  for query in queries {
    let mut known = false
    for parent in parents {
      if parent == query.parent_id {
        known = true
        break
      }
    }
    if !known {
      if tree.node(query.node_id) is None {
        return Err(InvalidTree)
      }
      parents.push(query.parent_id)
      contexts.push(query.context)
    }
  }
  Ok((parents, contexts))
}

///|
/// Turn a batch result into node-addressed probability rows. Keeping node IDs
/// at this boundary prevents an out-of-order inference backend from silently
/// scoring the wrong candidate.
fn append_batch_scores(
  tree : DraftTree,
  queries : Array[TreeQuery],
  parents : Array[Int?],
  logits_rows : Array[Array[Double]],
  scores : Array[NodeScore],
) -> Result[Unit, ModelTreeError] {
  if logits_rows.length() != parents.length() {
    return Err(BatchResultLengthMismatch)
  }
  for i in 0.. value
          None => return Err(InvalidTree)
        }
        expected_vocabulary = node.distribution.length()
        break
      }
    }
    if expected_vocabulary == 0 ||
      logits_rows[i].length() != expected_vocabulary {
      return Err(VocabularyMismatch)
    }
    let distribution = match softmax(logits_rows[i]) {
      Ok(value) => value
      Err(_) => return Err(InvalidLogits)
    }
    for query in queries {
      if query.parent_id == parents[i] {
        scores.push(NodeScore::new(query.node_id, distribution))
      }
    }
  }
  Ok(())
}

///|
/// Score every reachable parent context in one real batch-provider invocation.
/// The callback receives full token prefixes, so it can use a remote endpoint,
/// a tensor runtime, or a cache-aware native backend without changing the
/// verifier. Returned rows must retain input order.
pub fn score_tree_from_batch_model(
  tree : DraftTree,
  model : (Array[Array[Int]]) -> Result[Array[Array[Double]], String],
) -> Result[ScoredTree, ModelTreeError] {
  let plan = match plan_tree_batch(tree) {
    Ok(value) => value
    Err(_) => return Err(InvalidTree)
  }
  let queries = plan.flatten()
  let (parents, contexts) = match unique_parent_queries(tree, queries) {
    Ok(value) => value
    Err(error) => return Err(error)
  }
  let logits_rows = match model(contexts) {
    Ok(value) => value
    Err(error) => return Err(ProviderFailure(error))
  }
  let scores : Array[NodeScore] = []
  match append_batch_scores(tree, queries, parents, logits_rows, scores) {
    Ok(_) => ()
    Err(error) => return Err(error)
  }
  let distributions = match align_node_scores(tree, scores) {
    Ok(value) => value
    Err(_) => return Err(InvalidTree)
  }
  Ok({ distributions, provider_calls: 1 })
}

///|
/// Score one tree depth at a time. This is useful to backends that retain KV
/// cache state across levels or impose a maximum batch size. Unlike the
/// sequential adapter, siblings still share one batch row and `provider_calls`
/// is the number of actual batch invocations.
pub fn score_tree_layers_from_batch_model(
  tree : DraftTree,
  model : (Array[Array[Int]]) -> Result[Array[Array[Double]], String],
) -> Result[ScoredTree, ModelTreeError] {
  let plan = match plan_tree_batch(tree) {
    Ok(value) => value
    Err(_) => return Err(InvalidTree)
  }
  let scores : Array[NodeScore] = []
  let mut calls = 0
  for depth in 1..<(plan.layer_count() + 1) {
    let queries = plan.layer(depth)
    let (parents, contexts) = match unique_parent_queries(tree, queries) {
      Ok(value) => value
      Err(error) => return Err(error)
    }
    let logits_rows = match model(contexts) {
      Ok(value) => value
      Err(error) => return Err(ProviderFailure(error))
    }
    match append_batch_scores(tree, queries, parents, logits_rows, scores) {
      Ok(_) => ()
      Err(error) => return Err(error)
    }
    calls = calls + 1
  }
  let distributions = match align_node_scores(tree, scores) {
    Ok(value) => value
    Err(_) => return Err(InvalidTree)
  }
  Ok({ distributions, provider_calls: calls })
}