///| Preflight diagnostics for adapters that integrate TreeSpec with a model

///| runtime. They turn common malformed-fixture failures into a short report

///|
/// suitable for CI logs or a command-line health check.
pub enum DiagnosticStatus {
  Passed
  Failed
} derive(Eq, Debug)

///|
pub fn DiagnosticStatus::label(self : DiagnosticStatus) -> String {
  match self {
    Passed => "passed"
    Failed => "failed"
  }
}

///|
/// One named, independently meaningful preflight check.
pub struct DiagnosticCheck {
  name : String
  status : DiagnosticStatus
  detail : String
}

///|
pub fn DiagnosticCheck::passed(self : DiagnosticCheck) -> Bool {
  self.status == Passed
}

///|
pub fn DiagnosticCheck::render(self : DiagnosticCheck) -> String {
  self.status.label() + " " + self.name + ": " + self.detail
}

///| A report never throws: failed checks are data so callers can display all

///|
/// actionable issues at once instead of fixing one exception at a time.
pub struct DiagnosticReport {
  checks : Array[DiagnosticCheck]
}

///|
pub fn DiagnosticReport::empty() -> DiagnosticReport {
  { checks: [] }
}

///|
pub fn DiagnosticReport::add(
  self : DiagnosticReport,
  name : String,
  passed : Bool,
  detail : String,
) -> Unit {
  self.checks.push({
    name,
    status: if passed {
      Passed
    } else {
      Failed
    },
    detail,
  })
}

///|
pub fn DiagnosticReport::passed(self : DiagnosticReport) -> Bool {
  for check in self.checks {
    if !check.passed() {
      return false
    }
  }
  true
}

///|
pub fn DiagnosticReport::failure_count(self : DiagnosticReport) -> Int {
  let mut count = 0
  for check in self.checks {
    if !check.passed() {
      count = count + 1
    }
  }
  count
}

///|
pub fn DiagnosticReport::render(self : DiagnosticReport) -> String {
  let mut text = "TreeSpec preflight\n"
  for check in self.checks {
    text = text + "- " + check.render() + "\n"
  }
  text +
  "result=" +
  (if self.passed() { "passed" } else { "failed" }) +
  " failures=" +
  self.failure_count().to_string() +
  "\n"
}

///|
/// Check a single-path proposal without attempting target-model execution.
pub fn diagnose_proposal(proposal : DraftProposal) -> DiagnosticReport {
  let report = DiagnosticReport::empty()
  match validate_proposal(proposal) {
    Ok(_) =>
      report.add("proposal", true, "token ids and distributions are valid")
    Err(_) => report.add("proposal", false, "validation rejected the proposal")
  }
  report.add(
    "proposal-length",
    proposal.length() > 0,
    "candidate_tokens=" + proposal.length().to_string(),
  )
  report
}

///| Check structural tree validity and whether every node can be converted to

///|
/// a target-query context by the batch planner.
pub fn diagnose_tree(tree : DraftTree) -> DiagnosticReport {
  let report = DiagnosticReport::empty()
  match validate_tree(tree) {
    Ok(_) =>
      report.add("tree", true, "parent links and distributions are valid")
    Err(_) => report.add("tree", false, "validation rejected the tree")
  }
  match plan_tree_batch(tree) {
    Ok(plan) =>
      report.add(
        "tree-batch-plan",
        plan.query_count() == tree.node_count(),
        "query_count=" + plan.query_count().to_string(),
      )
    Err(_) =>
      report.add("tree-batch-plan", false, "batch planner rejected the tree")
  }
  report
}

///| Check a generated simulator workload before a benchmark run. This checks

///| row counts and vocabulary agreement without depending on a particular

///|
/// acceptance outcome.
pub fn diagnose_workload(schedule : Array[SimulationRound]) -> DiagnosticReport {
  let report = DiagnosticReport::empty()
  report.add(
    "workload-not-empty",
    schedule.length() > 0,
    "rounds=" + schedule.length().to_string(),
  )
  for index in 0.. 0 &&
      round.draft_logits.length() == round.target_logits.length() &&
      round.draft_logits.length() == round.accept_uniforms.length() &&
      round.draft_logits.length() == round.draft_uniforms.length() &&
      round.draft_logits.length() == round.fallback_uniforms.length()
    report.add(
      "workload-round-" + index.to_string(),
      shape_ok,
      "proposal_depth=" + round.draft_logits.length().to_string(),
    )
    if shape_ok {
      report.add(
        "workload-values-" + index.to_string(),
        run_round([], round) is Ok(_),
        "probability, vocabulary and random-input validation",
      )
    }
  }
  report
}

///|
/// Compare two provider outputs with a small numerical tolerance. Benchmark
/// callers need a stable model callback: changing logits for the same context
/// makes a matched baseline/adaptive comparison uninterpretable.
fn same_logits(
  left : Array[Array[Double]],
  right : Array[Array[Double]],
) -> Bool {
  if left.length() != right.length() {
    return false
  }
  for row_index in 0.. 1.0e-12 {
        return false
      }
    }
  }
  true
}

///|
/// Call a batch provider twice on fixed token contexts and report whether it
/// satisfies TreeSpec's integration contract. This preflight performs no
/// decoding; use it before a real benchmark to catch row reordering, a wrong
/// vocabulary, invalid logits or a stateful provider.
pub fn diagnose_batch_provider(
  contexts : Array[Array[Int]],
  expected_vocabulary : Int,
  model : (Array[Array[Int]]) -> Result[Array[Array[Double]], String],
) -> DiagnosticReport {
  let report = DiagnosticReport::empty()
  let request_ok = contexts.length() > 0 && expected_vocabulary > 0
  report.add(
    "batch-request",
    request_ok,
    "contexts=" +
    contexts.length().to_string() +
    " expected_vocabulary=" +
    expected_vocabulary.to_string(),
  )
  if !request_ok {
    return report
  }
  let first = match model(contexts.copy()) {
    Ok(value) => value
    Err(error) => {
      report.add("batch-provider", false, error)
      return report
    }
  }
  let row_count_ok = first.length() == contexts.length()
  report.add(
    "batch-row-count",
    row_count_ok,
    "received=" + first.length().to_string(),
  )
  if !row_count_ok {
    return report
  }
  let mut vocabulary_ok = true
  let mut logits_ok = true
  for row in first {
    if row.length() != expected_vocabulary {
      vocabulary_ok = false
    }
    if softmax(row) is Err(_) {
      logits_ok = false
    }
  }
  report.add("batch-vocabulary", vocabulary_ok, "all rows have expected width")
  report.add("batch-logits", logits_ok, "all rows produce valid softmax values")
  if !vocabulary_ok || !logits_ok {
    return report
  }
  let second = match model(contexts.copy()) {
    Ok(value) => value
    Err(error) => {
      report.add("batch-repeat", false, error)
      return report
    }
  }
  report.add(
    "batch-determinism",
    same_logits(first, second),
    "same contexts return equal logits within 1e-12",
  )
  report
}