///| Command implementation for the small offline executable. Keeping command

///| behavior in the library makes it testable without spawning a process and

///|
/// lets downstream tools reuse the same demonstrations.
pub enum CliError {
  UnknownCommand(String)
  DemoFailure
  WorkloadFailure
  TreePlanFailure
  ExperimentFailure
} derive(Eq, Debug)

///|
pub fn cli_help() -> String {
  "MoonBit TreeSpec\n\n" +
  "Usage: moon run cmd/main -- [command]\n\n" +
  "Commands:\n" +
  "  demo       Run the fixed two-round verification example (default).\n" +
  "  workload   Run a seeded synthetic draft/target workload.\n" +
  "  tree-plan  Print target-query contexts for a branching draft tree.\n" +
  "  experiment Compare context-dependent tree decoding with a matched baseline.\n" +
  "  batch-experiment Run context-sensitive decoding through batch callbacks.\n" +
  "  benchmark  Report matched baseline/adaptive request work.\n" +
  "  preflight  Validate a batch-model Provider contract.\n" +
  "  help       Print this message.\n"
}

///|
pub fn run_workload_demo() -> Result[String, CliError] {
  let config = WorkloadConfig::demo()
  let seed = 20260904
  let schedule = match generate_workload(config, seed) {
    Ok(value) => value
    Err(_) => return Err(WorkloadFailure)
  }
  let result = match simulate([1, 2], schedule) {
    Ok(value) => value
    Err(_) => return Err(WorkloadFailure)
  }
  Ok(describe_workload(config, seed) + "\n" + render_simulation(result))
}

///| A runnable branching example showing how an adapter receives one context

///| per tree node. It deliberately prints ids and contexts rather than logits,

///|
/// which keeps terminal output short and deterministic.
pub fn run_tree_plan_demo() -> Result[String, CliError] {
  let tree = match
    make_tree([101, 102], [
      {
        id: 0,
        parent: None,
        token: 7,
        distribution: [0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.3],
        depth: 1,
      },
      {
        id: 1,
        parent: None,
        token: 3,
        distribution: [0.1, 0.1, 0.1, 0.4, 0.1, 0.1, 0.05, 0.05],
        depth: 1,
      },
      {
        id: 2,
        parent: Some(0),
        token: 2,
        distribution: [0.1, 0.1, 0.4, 0.1, 0.1, 0.1, 0.05, 0.05],
        depth: 2,
      },
      {
        id: 3,
        parent: Some(0),
        token: 5,
        distribution: [0.1, 0.1, 0.1, 0.1, 0.1, 0.4, 0.05, 0.05],
        depth: 2,
      },
    ]) {
    Ok(value) => value
    Err(_) => return Err(TreePlanFailure)
  }
  match plan_tree_batch(tree) {
    Ok(plan) => Ok(plan.render())
    Err(_) => Err(TreePlanFailure)
  }
}

///| A deterministic stand-in for a host model provider. The command exercises

///|
/// the same benchmark API users call with a real batch inference callback.
pub fn run_benchmark_demo() -> Result[String, CliError] {
  let policy = match AdaptiveTreePolicy::new(1, 2, 1, 3, 14, 0.3, 0.65) {
    Ok(value) => value
    Err(_) => return Err(ExperimentFailure)
  }
  let target = fn(
    contexts : Array[Array[Int]],
  ) -> Result[Array[Array[Double]], String] {
    let rows : Array[Array[Double]] = []
    for context in contexts {
      let last = if context.is_empty() {
        0
      } else {
        context[context.length() - 1]
      }
      rows.push(if last == 0 { [0.0, 2.0, -1.0] } else { [1.0, -1.0, 0.5] })
    }
    Ok(rows)
  }
  let draft = fn(
    contexts : Array[Array[Int]],
  ) -> Result[Array[Array[Double]], String] {
    let rows : Array[Array[Double]] = []
    for context in contexts {
      let last = if context.is_empty() {
        0
      } else {
        context[context.length() - 1]
      }
      rows.push(if last == 0 { [0.0, 1.0, -0.5] } else { [0.8, 0.0, 0.3] })
    }
    Ok(rows)
  }
  match
    benchmark_adaptive_batch_decoding(
      "synthetic-context-model",
      [0],
      draft,
      target,
      policy,
      32,
      20260905,
    ) {
    Ok(value) => Ok(value.render())
    Err(_) => Err(ExperimentFailure)
  }
}

///| Demonstrate the provider contract that an external model adapter must

///|
/// satisfy before it is used for decoding or benchmark collection.
pub fn run_provider_preflight_demo() -> Result[String, CliError] {
  let report = diagnose_batch_provider([[0], [0, 1]], 3, contexts => {
    let rows : Array[Array[Double]] = []
    for context in contexts {
      rows.push(
        if context.length() == 1 {
          [0.0, 1.0, -1.0]
        } else {
          [1.0, 0.0, -1.0]
        },
      )
    }
    Ok(rows)
  })
  if report.passed() {
    Ok(report.render())
  } else {
    Err(ExperimentFailure)
  }
}

///| Dispatch a command without any environment dependency. Unknown commands

///|
/// return a typed error so embedding applications can show their own help.
pub fn run_cli_command(command : String) -> Result[String, CliError] {
  match command {
    "demo" =>
      match run_demo() {
        Ok(text) => Ok(text)
        Err(_) => Err(DemoFailure)
      }
    "workload" => run_workload_demo()
    "tree-plan" => run_tree_plan_demo()
    "experiment" =>
      match run_tree_experiment_demo() {
        Ok(value) => Ok(value)
        Err(_) => Err(ExperimentFailure)
      }
    "batch-experiment" =>
      match run_tree_batch_experiment_demo() {
        Ok(value) => Ok(value)
        Err(_) => Err(ExperimentFailure)
      }
    "benchmark" => run_benchmark_demo()
    "preflight" => run_provider_preflight_demo()
    "help" => Ok(cli_help())
    _ => Err(UnknownCommand(command))
  }
}