use anyhow::Result;
use chrono::Utc;
use futures::StreamExt;
use std::collections::HashMap;
use std::path::PathBuf;
use tracing::info;
use wme_client::RealtimeConnectOptions;
use crate::client::map_client_error;
use crate::commands::GlobalOpts;
use crate::output::{format_table, OutputHandler};
use crate::util::build_params;
pub async fn stream(
since: Option<String>,
parts: Vec<u32>,
offsets: Option<String>,
since_per_partition: Option<String>,
ndjson: bool,
opts: &GlobalOpts,
) -> Result<()> {
info!("Connecting to realtime stream");
println!("Connecting to realtime stream...");
println!("Press Ctrl+C to stop\n");
let client = opts.build_client().await?;
let mut options = RealtimeConnectOptions::default();
if let Some(since_str) = since {
let since_dt = chrono::DateTime::parse_from_rfc3339(&since_str)
.map_err(|e| anyhow::anyhow!("Invalid since timestamp '{}': {}", since_str, e))?;
options.since = Some(since_dt.with_timezone(&Utc));
}
if !parts.is_empty() {
options.partitions = Some(parts);
}
if let Some(offsets_str) = offsets {
let offsets_map: HashMap<String, u64> = serde_json::from_str(&offsets_str)
.map_err(|e| anyhow::anyhow!("Invalid offsets JSON '{}': {}", offsets_str, e))?;
options.offsets = Some(offsets_map);
}
if let Some(since_per_part_str) = since_per_partition {
let since_map: HashMap<String, String> = serde_json::from_str(&since_per_part_str)
.map_err(|e| {
anyhow::anyhow!(
"Invalid since_per_partition JSON '{}': {}",
since_per_part_str,
e
)
})?;
let mut since_per_part: HashMap<String, chrono::DateTime<Utc>> = HashMap::new();
for (k, v) in since_map {
let dt = chrono::DateTime::parse_from_rfc3339(&v).map_err(|e| {
anyhow::anyhow!("Invalid timestamp in since_per_partition '{}': {}", v, e)
})?;
since_per_part.insert(k, dt.with_timezone(&Utc));
}
options.since_per_partition = Some(since_per_part);
}
let params = build_params(opts)?;
if ndjson {
println!("NDJSON format requested - streaming to stdout...");
println!("(Note: NDJSON streaming via client library is not yet fully implemented)");
return Ok(());
}
let mut stream = client
.realtime()
.connect(&options, params.as_ref())
.await
.map_err(map_client_error)?;
while let Some(result) = stream.next().await {
match result {
Ok(update) => {
let json = serde_json::to_string(&update)?;
println!("{}", json);
}
Err(e) => {
eprintln!("Error in stream: {}", e);
}
}
}
Ok(())
}
pub async fn batches_list(date: &str, hour: &str, opts: &GlobalOpts) -> Result<()> {
info!("Listing batches for {} {}", date, hour);
let client = opts.build_client().await?;
let params = build_params(opts)?;
let batches = client
.realtime()
.list_batches_with_params(date, hour, params.as_ref())
.await
.map_err(map_client_error)?;
let output = OutputHandler::new(opts.output);
if matches!(opts.output, crate::output::OutputFormat::Table) {
let headers = vec!["Identifier", "Name", "Size"];
let rows: Vec<Vec<String>> = batches
.iter()
.map(|batch| {
let size_str = crate::util::format_api_size(&batch.size);
vec![batch.identifier.clone(), batch.name.clone(), size_str]
})
.collect();
output.output_raw(&format_table(&headers, &rows));
} else {
let json_value = serde_json::to_value(&batches)?;
output.output(&json_value)?;
}
Ok(())
}
pub async fn batches_info(
date: &str,
hour: &str,
identifier: &str,
opts: &GlobalOpts,
) -> Result<()> {
info!("Getting batch info: {} for {} {}", identifier, date, hour);
let client = opts.build_client().await?;
let params = build_params(opts)?;
let batch = client
.realtime()
.get_batch_info_with_params(date, hour, identifier, params.as_ref())
.await
.map_err(map_client_error)?;
let output = OutputHandler::new(opts.output);
let json_value = serde_json::to_value(&batch)?;
output.output(&json_value)?;
Ok(())
}
pub async fn batches_download(
_date: &str,
_hour: &str,
identifier: &str,
output_file: Option<PathBuf>,
range: Option<String>,
opts: &GlobalOpts,
) -> Result<()> {
let output_path =
output_file.unwrap_or_else(|| PathBuf::from(format!("{}.tar.gz", identifier)));
info!(
"Downloading batch: {} to {}",
identifier,
output_path.display()
);
println!("Downloading batch: {}", identifier);
println!("Output: {}", output_path.display());
let client = opts.build_client().await?;
let date = _date;
let hour = _hour;
let range_ref = range.as_deref();
let mut stream = client
.realtime()
.download_batch(date, hour, identifier, range_ref)
.await
.map_err(map_client_error)?;
let mut file = tokio::fs::File::create(&output_path).await?;
while let Some(chunk_result) = stream.next().await {
let chunk = chunk_result.map_err(map_client_error)?;
tokio::io::AsyncWriteExt::write_all(&mut file, &chunk).await?;
}
let metadata = tokio::fs::metadata(&output_path).await?;
let size = metadata.len();
let size_str = crate::util::format_size(size);
println!("✓ Download complete!");
println!(" File: {}", output_path.display());
println!(" Size: {}", size_str);
Ok(())
}