use anyhow::{Context, Result};
use clap::{Parser, ValueEnum};
use bee_check::{
DEFAULT_GATEWAY, OutputFormat, ParsedInput, ReseedRequest, check_gateways,
check_multi_vantage, check_stamp, drill_down, merge_gateways, parse_input, render_report,
render_stamp_status, reseed, resolve_feed,
};
#[derive(Parser, Debug)]
#[command(
name = "bee-check",
version,
about = "Retrievability checker for Swarm references",
long_about = "Probe one or more Bee nodes to determine whether a Swarm \
reference is retrievable from the network. Supports \
per-chunk drill-down and one-shot re-seed via stewardship."
)]
struct Cli {
#[arg(value_name = "INPUT")]
reference: String,
#[arg(short = 'b', long = "bee", value_name = "URL")]
bee: Vec<String>,
#[arg(long = "gateway", value_name = "URL", conflicts_with = "no_gateway")]
gateway: Vec<String>,
#[arg(long)]
no_gateway: bool,
#[arg(long)]
per_chunk: bool,
#[arg(long)]
reseed: bool,
#[arg(long, value_name = "ID", requires = "reseed")]
stamp: Option<String>,
#[arg(long, default_value_t = 60)]
timeout: u64,
#[arg(long, default_value_t = 8)]
concurrency: usize,
#[arg(long, value_enum, default_value_t = OutputKind::Text)]
output: OutputKind,
}
#[derive(Copy, Clone, Debug, ValueEnum)]
enum OutputKind {
Text,
Json,
}
impl From<OutputKind> for OutputFormat {
fn from(k: OutputKind) -> Self {
match k {
OutputKind::Text => OutputFormat::Text,
OutputKind::Json => OutputFormat::Json,
}
}
}
fn default_bees() -> Vec<String> {
std::env::var("BEE_API_URL")
.ok()
.into_iter()
.chain(std::iter::once("http://localhost:1633".to_string()))
.take(1)
.collect()
}
#[tokio::main]
async fn main() -> Result<()> {
let cli = Cli::parse();
let bees: Vec<String> = if cli.bee.is_empty() {
default_bees()
} else {
cli.bee.clone()
};
let timeout = std::time::Duration::from_secs(cli.timeout);
let (reference, resolution) = match parse_input(&cli.reference) {
ParsedInput::Reference(r) => (r, None),
ParsedInput::Feed { owner, topic } => {
let first_bee = bees
.first()
.context("feed resolution requires at least one --bee URL")?;
let (r, res) = resolve_feed(first_bee, &owner, &topic, timeout)
.await
.context("feed resolution failed")?;
eprintln!("resolved feed -> {r}");
(r, Some(res))
}
};
let gateways: Vec<String> = if cli.no_gateway {
Vec::new()
} else if cli.gateway.is_empty() {
vec![DEFAULT_GATEWAY.to_string()]
} else {
cli.gateway.clone()
};
let (vantage_report, gateway_results) = tokio::join!(
check_multi_vantage(&reference, &bees, timeout),
check_gateways(&reference, &gateways, timeout),
);
let mut report = vantage_report.context("multi-vantage check failed")?;
let gateways_out = gateway_results.context("gateway probe failed")?;
report = merge_gateways(report, gateways_out);
report.resolution = resolution;
let report = if cli.per_chunk {
drill_down(report, &bees, timeout, cli.concurrency)
.await
.context("per-chunk drill-down failed")?
} else {
report
};
print!("{}", render_report(&report, cli.output.into()));
if cli.reseed {
let stamp = cli
.stamp
.as_ref()
.context("--reseed requires --stamp <batch-id>")?;
let target_bee = bees.first().context("no bee URL for --reseed")?;
let status = check_stamp(target_bee, stamp, timeout)
.await
.context("stamp pre-flight failed")?;
eprint!("{}", render_stamp_status(&status));
if !status.exists || !status.usable {
anyhow::bail!("refusing to re-seed: stamp is not usable; see warnings above");
}
let req = ReseedRequest {
reference: reference.clone(),
bee_url: target_bee.clone(),
batch_id: stamp.clone(),
timeout,
};
reseed(req).await.context("re-seed failed")?;
eprintln!("re-seeded {} via {}", reference, target_bee);
}
let any_retrievable = report.vantages.iter().any(|v| v.retrievable == Some(true))
|| report.gateways.iter().any(|g| g.retrievable == Some(true));
if !any_retrievable {
std::process::exit(2);
}
Ok(())
}