///| End-to-end speculative decoding using two `ReplayModel` instances. This

///| is deliberately a reference adapter: it demonstrates every context query

///|
/// that a real draft/target inference integration must make.
pub enum ReplayDecodeError {
  InvalidDepth(Int)
  DraftModelFailure(Int)
  TargetModelFailure(Int)
  ProbabilityFailure(Int)
  VerificationFailure
} derive(Eq, Debug)

///| A completed replay round includes the proposal and verification outcome so

///| callers can inspect accepted prefixes instead of treating decoding as a

///|
/// black box.
pub struct ReplayRoundResult {
  proposal : DraftProposal
  verification : VerifyResult
}

///|
pub fn ReplayRoundResult::proposal(self : ReplayRoundResult) -> DraftProposal {
  self.proposal
}

///|
pub fn ReplayRoundResult::verification(
  self : ReplayRoundResult,
) -> VerifyResult {
  self.verification
}

///|
fn copy_tokens(tokens : Array[Int]) -> Array[Int] {
  let copied : Array[Int] = []
  for token in tokens {
    copied.push(token)
  }
  copied
}

///|
/// Sample every draft token from the distribution used for verification.
/// The supplied draws must be independent of acceptance and fallback draws.
pub fn proposal_from_replay(
  draft : ReplayModel,
  prefix : Array[Int],
  depth : Int,
  draft_uniforms : Array[Double],
) -> Result[DraftProposal, ReplayDecodeError] {
  if depth <= 0 {
    return Err(InvalidDepth(depth))
  }
  if draft_uniforms.length() != depth {
    return Err(ProbabilityFailure(0))
  }
  for i in 0.. value
      Err(_) => return Err(DraftModelFailure(index))
    }
    let distribution = match softmax(logits) {
      Ok(value) => value
      Err(_) => return Err(ProbabilityFailure(index))
    }
    let token = match sample_categorical(distribution, draft_uniforms[index]) {
      Ok(value) => value
      Err(_) => return Err(ProbabilityFailure(index))
    }
    tokens.push(token)
    distributions.push(distribution)
    context.push(token)
  }
  match make_proposal(prefix, tokens, distributions) {
    Ok(proposal) => Ok(proposal)
    Err(_) => Err(ProbabilityFailure(depth))
  }
}

///| Query target logits for every proposed prefix, then run the exact

///| single-path verifier. The target is queried before verification because a

///| batched production target model would return the same set of positions in

///|
/// one forward pass.
pub fn verify_from_replay(
  target : ReplayModel,
  proposal : DraftProposal,
  accept_uniforms : Array[Double],
  fallback_uniforms : Array[Double],
) -> Result[VerifyResult, ReplayDecodeError] {
  let context = copy_tokens(proposal.prefix)
  let target_distributions : Array[Array[Double]] = []
  for index in 0.. value
      Err(_) => return Err(TargetModelFailure(index))
    }
    let distribution = match softmax(logits) {
      Ok(value) => value
      Err(_) => return Err(ProbabilityFailure(index))
    }
    target_distributions.push(distribution)
    context.push(proposal.tokens[index].token)
  }
  match
    verify_proposal(
      proposal, target_distributions, accept_uniforms, fallback_uniforms,
    ) {
    Ok(result) => Ok(result)
    Err(_) => Err(VerificationFailure)
  }
}

///|
/// Execute one replay-backed speculative round.
pub fn run_replay_round(
  draft : ReplayModel,
  target : ReplayModel,
  prefix : Array[Int],
  depth : Int,
  accept_uniforms : Array[Double],
  fallback_uniforms : Array[Double],
  draft_uniforms : Array[Double],
) -> Result[ReplayRoundResult, ReplayDecodeError] {
  let proposal = match
    proposal_from_replay(draft, prefix, depth, draft_uniforms) {
    Ok(value) => value
    Err(error) => return Err(error)
  }
  let verification = match
    verify_from_replay(target, proposal, accept_uniforms, fallback_uniforms) {
    Ok(value) => value
    Err(error) => return Err(error)
  }
  Ok({ proposal, verification })
}

///| Append one replay round to a decoding history and account for target work.

///|
/// This helper makes it hard for benchmark callers to forget metric updates.
pub fn append_replay_round(
  generated : Array[Int],
  metrics : DecodeMetrics,
  result : ReplayRoundResult,
) -> Unit {
  metrics.record(result.verification, result.proposal)
  for token in result.verification.emitted {
    generated.push(token)
  }
}