wme-cli 0.1.3

CLI tool for the Wikimedia Enterprise API
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;

/// Stream realtime updates
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?;

    // Build RealtimeConnectOptions
    let mut options = RealtimeConnectOptions::default();

    // Parse since timestamp
    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));
    }

    // Set partitions
    if !parts.is_empty() {
        options.partitions = Some(parts);
    }

    // Parse offsets JSON
    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);
    }

    // Parse since_per_partition JSON
    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);
    }

    // Build request params (for filters)
    let params = build_params(opts)?;

    if ndjson {
        // For NDJSON, we need to use a different approach
        // For now, just print a message that this is not yet fully implemented
        println!("NDJSON format requested - streaming to stdout...");
        println!("(Note: NDJSON streaming via client library is not yet fully implemented)");
        return Ok(());
    }

    // Connect to SSE stream
    let mut stream = client
        .realtime()
        .connect(&options, params.as_ref())
        .await
        .map_err(map_client_error)?;

    // Stream updates
    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(())
}

/// List batches for a date/hour
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);

    // Format as table if requested
    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(())
}

/// Get batch metadata
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(())
}

/// Download a batch
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?;

    // Note: The client method needs date and hour, but the CLI args have them
    // For now, we'll use placeholder values - this needs to be fixed in the API
    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)?;

    // Stream to file
    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?;
    }

    // Get final size
    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(())
}