///| Token constraints are applied after sampling filters and before a token is

///| drawn. They model common serving requirements such as a JSON-only grammar

///| adapter, tool-call vocabulary restriction, or a caller-provided EOS rule

///|
/// without coupling TreeSpec to one tokenizer.
pub enum ConstraintError {
  InvalidVocabulary(Int)
  InvalidToken(Int)
  EmptyAllowedSet
  EmptyStopSequence
  DistributionMismatch
  SamplingFailure
} derive(Eq, Debug)

///|
/// A vocabulary-sized boolean mask. `true` means the token may be emitted.
pub struct TokenMask {
  allowed : Array[Bool]
}

///|
pub fn TokenMask::allow_all(
  vocabulary_size : Int,
) -> Result[TokenMask, ConstraintError] {
  if vocabulary_size <= 0 {
    return Err(InvalidVocabulary(vocabulary_size))
  }
  let allowed : Array[Bool] = []
  for _ in 0.. Result[TokenMask, ConstraintError] {
  if vocabulary_size <= 0 {
    return Err(InvalidVocabulary(vocabulary_size))
  }
  let allowed : Array[Bool] = []
  for _ in 0..= vocabulary_size {
      return Err(InvalidToken(token))
    }
    allowed[token] = true
  }
  let mut any = false
  for value in allowed {
    if value {
      any = true
    }
  }
  if !any {
    Err(EmptyAllowedSet)
  } else {
    Ok({ allowed, })
  }
}

///|
/// Start with every token allowed and selectively prohibit listed ids.
pub fn TokenMask::except(
  vocabulary_size : Int,
  token_ids : Array[Int],
) -> Result[TokenMask, ConstraintError] {
  let mask = match TokenMask::allow_all(vocabulary_size) {
    Ok(value) => value
    Err(error) => return Err(error)
  }
  for token in token_ids {
    if token < 0 || token >= vocabulary_size {
      return Err(InvalidToken(token))
    }
    mask.allowed[token] = false
  }
  let mut any = false
  for value in mask.allowed {
    if value {
      any = true
    }
  }
  if !any {
    Err(EmptyAllowedSet)
  } else {
    Ok(mask)
  }
}

///|
pub fn TokenMask::vocabulary_size(self : TokenMask) -> Int {
  self.allowed.length()
}

///|
pub fn TokenMask::allows(self : TokenMask, token : Int) -> Bool {
  token >= 0 && token < self.allowed.length() && self.allowed[token]
}

///| Intersect two independent constraints, preserving neither caller's mutable

///| storage. This is useful when a grammar mask and an application policy are

///|
/// both active for one decoding step.
pub fn TokenMask::intersect(
  self : TokenMask,
  other : TokenMask,
) -> Result[TokenMask, ConstraintError] {
  if self.allowed.length() != other.allowed.length() {
    return Err(DistributionMismatch)
  }
  let allowed : Array[Bool] = []
  let mut any = false
  for index in 0.. Result[Array[Double], ConstraintError] {
  if distribution.length() != self.allowed.length() {
    return Err(DistributionMismatch)
  }
  let weights : Array[Double] = []
  for index in 0.. Ok(value)
    Err(_) => Err(EmptyAllowedSet)
  }
}

///| A token-id stop sequence. It is deliberately tokenizer-neutral: a caller

///|
/// can represent EOS, a delimiter, or a multi-token marker with the same API.
pub struct StopSequence {
  tokens : Array[Int]
}

///|
pub fn StopSequence::new(
  tokens : Array[Int],
) -> Result[StopSequence, ConstraintError] {
  if tokens.length() == 0 {
    return Err(EmptyStopSequence)
  }
  for token in tokens {
    if token < 0 {
      return Err(InvalidToken(token))
    }
  }
  Ok({ tokens, })
}

///|
pub fn StopSequence::length(self : StopSequence) -> Int {
  self.tokens.length()
}

///|
/// True only if the current output ends in this exact sequence.
pub fn StopSequence::matches_suffix(
  self : StopSequence,
  output : Array[Int],
) -> Bool {
  if output.length() < self.tokens.length() {
    return false
  }
  let offset = output.length() - self.tokens.length()
  for index in 0.. Result[Int, ConstraintError] {
  let distribution = match distribution_for_sampling(logits, config) {
    Ok(value) => value
    Err(_) => return Err(SamplingFailure)
  }
  let constrained = match mask.apply(distribution) {
    Ok(value) => value
    Err(error) => return Err(error)
  }
  match sample_categorical(constrained, unit_interval) {
    Ok(token) => Ok(token)
    Err(_) => Err(SamplingFailure)
  }
}