knut-bifrost 0.1.3

//! `bifrost` — drive and observe the file → graph lifecycle headlessly.
//!
//! This is the CLI front door to the [`Listener`](knut_bifrost::Listener): it
//! watches a directory of file drops, runs them through a mapping into a graph
//! sink, and emits **structured tracing events** at every stage — `file_detected`,
//! `files_read`, `table_synced`, `sink_applied`, `round_complete`. With
//! `--json-logs` those become one JSON object per line, so a test or bencher
//! can drive this binary and assert on the lifecycle with no human in the loop.
//!
//! `--dry-run` swaps the FalkorDB sink for a counting null sink, so the
//! detect→transform half of the pipeline can be exercised with no database.
//!
//! ```text
//! bifrost run --mapping m.yml --files ./drops --graph demo \
//!     --falkor redis://127.0.0.1:6379 --once --json-logs
//! bifrost run --mapping m.yml --files ./drops --dry-run --once   # no DB
//! ```

use std::path::PathBuf;
use std::process::ExitCode;
use std::time::Duration;

use async_trait::async_trait;
use clap::{Parser, Subcommand};
use knut_bifrost::{
    Backoff, GraphMutation, Listener, MappingSpec, Result as BifrostResult, SinkReport, SyncState,
    TransformRegistry,
    sink::{FalkorDbSink, GraphSink},
    source::FileSource,
};

#[derive(Parser)]
#[command(
    name = "bifrost",
    about = "Drive the file → graph lifecycle with structured logs",
    version
)]
struct Cli {
    /// Emit one JSON object per log line (machine-readable) instead of pretty text.
    #[arg(long, global = true)]
    json_logs: bool,
    #[command(subcommand)]
    command: Command,
}

#[derive(Subcommand)]
// `Run` carries the whole CLI surface (mapping/files/falkor/backoff/metrics/…)
// and is constructed exactly once at startup, so the size gap to `GenSample`
// is intentional — boxing a parsed-once clap variant buys nothing.
#[allow(clippy::large_enum_variant)]
enum Command {
    /// Watch a directory of file drops and materialise them into a graph.
    Run {
        /// Mapping YAML (table → pipeline). Table names map to `<files>/<table>/`.
        #[arg(long)]
        mapping: PathBuf,
        /// Root directory of file drops to watch.
        #[arg(long)]
        files: PathBuf,
        /// FalkorDB / Redis URL (ignored with --dry-run).
        #[arg(long, default_value = "redis://127.0.0.1:6379")]
        falkor: String,
        /// Target graph name.
        #[arg(long, default_value = "bifrost")]
        graph: String,
        /// Sync-state ledger path (default: <files>/.bifrost-state.json).
        #[arg(long)]
        state: Option<PathBuf>,
        /// Use a **Valkey**-backed shared ledger at this URL instead of a file
        /// (so multiple instances share progress). Valkey 9+ / Redis-API.
        #[arg(long, value_name = "URL")]
        state_valkey: Option<String>,
        /// Valkey hash key for the ledger (with --state-valkey).
        #[arg(long, default_value = "bifrost:state")]
        state_key: String,
        /// Run exactly one sync round and exit (else poll forever).
        #[arg(long)]
        once: bool,
        /// Poll interval in seconds when looping.
        #[arg(long, default_value_t = 5)]
        interval: u64,
        /// Cap (seconds) for the failure backoff. A fully-failed round doubles
        /// the inter-round sleep up to this, resetting on the first round that
        /// makes progress. Defaults to 10× `--interval`.
        #[arg(long)]
        max_backoff: Option<u64>,
        /// On shutdown, write the cumulative run metrics as Prometheus text to
        /// this file (a `/metrics` endpoint can serve the same later).
        #[arg(long, value_name = "PATH")]
        metrics_file: Option<PathBuf>,
        /// Serve the cumulative run metrics live over HTTP at `/metrics` on this
        /// `ip:port` (e.g. `127.0.0.1:9090`) so a Prometheus scraper can pull
        /// them while bifröst runs — published after every round. Absent → no
        /// server (today's behavior). Requires building with the `metrics-http`
        /// feature.
        #[arg(long, value_name = "ADDR")]
        metrics_addr: Option<String>,
        /// Skip the database: count the graph mutation instead of writing it.
        #[arg(long)]
        dry_run: bool,
        /// Airflow API base URL — after each round, emit an AIP-82 asset event
        /// per ingested table to `{url}/api/v2/assets/events`.
        #[arg(long, value_name = "URL")]
        emit_airflow: Option<String>,
        /// Asset URI prefix for emitted events (event uri = `{prefix}{table}`).
        #[arg(long, default_value = "graph://")]
        asset_prefix: String,
        /// Bearer token for the Airflow API, if required.
        #[arg(long, value_name = "TOKEN")]
        airflow_token: Option<String>,
    },
    /// Write a sample parquet drop (`id`, `name`, `score`) under
    /// `<files>/<table>/<chunk>.parquet` — handy for trying `run` with no
    /// other tooling.
    GenSample {
        /// Root drops directory.
        #[arg(long)]
        files: PathBuf,
        /// Table (subdirectory) name.
        #[arg(long, default_value = "people")]
        table: String,
        /// Number of rows.
        #[arg(long, default_value_t = 1000)]
        rows: usize,
        /// Chunk file name.
        #[arg(long, default_value = "chunk0.parquet")]
        chunk: String,
    },
    /// Headless self-test: run the core transform→mutation path over a known
    /// in-memory table and report each invariant. Built `--features testmatrix`
    /// it also emits `functional` rows into the nornir test matrix. Exits
    /// non-zero if any check fails (so CI / the matrix sees a broken core).
    SelfTest,
}

fn main() -> ExitCode {
    let cli = Cli::parse();
    init_tracing(cli.json_logs);
    match run(cli.command) {
        Ok(()) => ExitCode::SUCCESS,
        Err(e) => {
            tracing::error!(target: "bifrost", error = %e, "fatal");
            ExitCode::FAILURE
        }
    }
}

/// Install a tracing subscriber: JSON lines or pretty text, level from
/// `RUST_LOG` (default `info`).
fn init_tracing(json: bool) {
    use tracing_subscriber::{EnvFilter, fmt};
    let filter = EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("info"));
    if json {
        fmt()
            .with_env_filter(filter)
            .json()
            .with_current_span(false)
            .init();
    } else {
        fmt().with_env_filter(filter).with_target(true).init();
    }
}

fn run(cmd: Command) -> Result<(), Box<dyn std::error::Error>> {
    match cmd {
        Command::SelfTest => {
            let ok = knut_bifrost::testmatrix::self_test();
            tracing::info!(target: "bifrost", ok, "self-test complete");
            if ok {
                Ok(())
            } else {
                Err("bifröst self-test failed (a core invariant did not hold)".into())
            }
        }
        Command::GenSample {
            files,
            table,
            rows,
            chunk,
        } => gen_sample(files, table, rows, chunk),
        Command::Run {
            mapping,
            files,
            falkor,
            graph,
            state,
            state_valkey,
            state_key,
            once,
            interval,
            max_backoff,
            metrics_file,
            metrics_addr,
            dry_run,
            emit_airflow,
            asset_prefix,
            airflow_token,
        } => run_pipeline(
            mapping,
            files,
            falkor,
            graph,
            state,
            state_valkey,
            state_key,
            once,
            interval,
            max_backoff,
            metrics_file,
            metrics_addr,
            dry_run,
            emit_airflow,
            asset_prefix,
            airflow_token,
        ),
    }
}

#[allow(clippy::too_many_arguments)]
fn run_pipeline(
    mapping: PathBuf,
    files: PathBuf,
    falkor: String,
    graph: String,
    state: Option<PathBuf>,
    state_valkey: Option<String>,
    state_key: String,
    once: bool,
    interval: u64,
    max_backoff: Option<u64>,
    metrics_file: Option<PathBuf>,
    metrics_addr: Option<String>,
    dry_run: bool,
    emit_airflow: Option<String>,
    asset_prefix: String,
    airflow_token: Option<String>,
) -> Result<(), Box<dyn std::error::Error>> {
    let spec = MappingSpec::load(&mapping)?;
    let state_path = state.unwrap_or_else(|| files.join(".bifrost-state.json"));
    let rt = tokio::runtime::Builder::new_multi_thread()
        .enable_all()
        .build()?;

    rt.block_on(async move {
        let source = FileSource::new(&files);
        // Valkey-backed shared ledger, or the default atomic file ledger.
        let sync_state = match &state_valkey {
            Some(url) => SyncState::open_valkey(url, &state_key).await?,
            None => SyncState::open(&state_path)?,
        };
        // AIP-82 emitter (closes the loop: ingested → asset event → DAG).
        let emitter: Option<(std::sync::Arc<dyn knut_bifrost::AssetEmitter>, String)> =
            emit_airflow.map(|url| {
                let e = knut_bifrost::airflow::AirflowEmitter::new(url, airflow_token);
                (
                    std::sync::Arc::new(e) as std::sync::Arc<dyn knut_bifrost::AssetEmitter>,
                    asset_prefix,
                )
            });
        tracing::info!(
            target: "bifrost",
            event = "start",
            files = %files.display(),
            graph = %graph,
            dry_run,
            once,
            emit_airflow = emitter.is_some(),
            metrics_http = metrics_addr.is_some(),
            state = if state_valkey.is_some() { "valkey" } else { "file" },
            "bifrost starting",
        );

        // Start the live `/metrics` HTTP server if `--metrics-addr` was given.
        // Held for the whole run so the endpoint stays up; the poll loop
        // publishes the cumulative metrics into its shared snapshot each round.
        // The `MetricsServer` is dropped (stopped + joined) when this scope ends.
        let metrics_server = start_metrics_server(metrics_addr.as_deref())?;

        if dry_run {
            let sink = NullSink;
            let mut l = Listener::new(spec, TransformRegistry::new(), source, sink, sync_state);
            if let Some((e, prefix)) = emitter {
                l = l.with_asset_emitter(e, prefix);
            }
            drive(
                l,
                once,
                interval,
                max_backoff,
                metrics_file,
                &metrics_server,
            )
            .await
        } else {
            let sink = FalkorDbSink::connect(&falkor, graph.clone()).await?;
            let mut l = Listener::new(spec, TransformRegistry::new(), source, sink, sync_state);
            if let Some((e, prefix)) = emitter {
                l = l.with_asset_emitter(e, prefix);
            }
            drive(
                l,
                once,
                interval,
                max_backoff,
                metrics_file,
                &metrics_server,
            )
            .await
        }
    })
}

// ── Live `/metrics` HTTP endpoint (feature `metrics-http`) ───────────────────
//
// The server start + per-round publish are abstracted behind `OptMetricsServer`
// and two helpers so the `drive` loop is identical whether or not the feature is
// compiled in. With the feature on, `OptMetricsServer` is an
// `Option<MetricsServer>` and `publish_metrics` forwards the cumulative metrics
// into the shared snapshot each round; without it, it's a zero-sized unit and
// both helpers are no-ops (and `--metrics-addr` warns that the build lacks the
// feature).

/// Carrier for the optional running metrics server, kept alive for the run.
#[cfg(feature = "metrics-http")]
type OptMetricsServer = Option<knut_bifrost::MetricsServer>;
#[cfg(not(feature = "metrics-http"))]
type OptMetricsServer = ();

/// Start the live `/metrics` HTTP server when `--metrics-addr` is given (and the
/// `metrics-http` feature is built in). A bind failure is fatal — the operator
/// asked for the endpoint, so failing loudly beats silently dropping it.
#[cfg(feature = "metrics-http")]
fn start_metrics_server(
    addr: Option<&str>,
) -> Result<OptMetricsServer, Box<dyn std::error::Error>> {
    match addr {
        Some(addr) => {
            let server = knut_bifrost::MetricsServer::start(addr)?;
            let bound = server
                .local_addr()
                .map(|a| a.to_string())
                .unwrap_or_else(|_| addr.into());
            tracing::info!(
                target: "metrics",
                event = "metrics_http_started",
                addr = %bound,
                "serving Prometheus metrics at http://{bound}/metrics",
            );
            Ok(Some(server))
        }
        None => Ok(None),
    }
}

/// Without the `metrics-http` feature there is no server; warn if the flag was
/// nonetheless supplied so the operator knows the endpoint isn't running.
#[cfg(not(feature = "metrics-http"))]
fn start_metrics_server(
    addr: Option<&str>,
) -> Result<OptMetricsServer, Box<dyn std::error::Error>> {
    if addr.is_some() {
        tracing::warn!(
            target: "metrics",
            event = "metrics_http_unavailable",
            "--metrics-addr was given but this build lacks the `metrics-http` feature; \
             no /metrics server started (rebuild with --features metrics-http)",
        );
    }
    Ok(())
}

/// Publish the listener's cumulative metrics into the live server's shared
/// snapshot so the next scrape renders them. No-op without the feature / server.
#[cfg(feature = "metrics-http")]
fn publish_metrics<Snk: GraphSink>(
    server: &OptMetricsServer,
    listener: &Listener<FileSource, Snk>,
) {
    if let Some(server) = server {
        server.handle().publish(listener.metrics());
    }
}

#[cfg(not(feature = "metrics-http"))]
fn publish_metrics<Snk: GraphSink>(
    _server: &OptMetricsServer,
    _listener: &Listener<FileSource, Snk>,
) {
}

/// Write a sample parquet drop so `bifrost run` can be tried with no other
/// tooling. Columns: `id: Int64`, `name: Utf8`, `score: Float64`.
fn gen_sample(
    files: PathBuf,
    table: String,
    rows: usize,
    chunk: String,
) -> Result<(), Box<dyn std::error::Error>> {
    use arrow_array::{Float64Array, Int64Array, RecordBatch, StringArray};
    use arrow_schema::{DataType, Field, Schema};
    use std::sync::Arc;

    let ids: Vec<i64> = (0..rows as i64).collect();
    let names: Vec<String> = (0..rows).map(|i| format!("person_{i}")).collect();
    let scores: Vec<f64> = (0..rows).map(|i| i as f64 * 0.5).collect();
    let batch = RecordBatch::try_new(
        Arc::new(Schema::new(vec![
            Field::new("id", DataType::Int64, false),
            Field::new("name", DataType::Utf8, false),
            Field::new("score", DataType::Float64, false),
        ])),
        vec![
            Arc::new(Int64Array::from(ids)),
            Arc::new(StringArray::from(names)),
            Arc::new(Float64Array::from(scores)),
        ],
    )?;

    let dir = files.join(&table);
    std::fs::create_dir_all(&dir)?;
    let path = dir.join(&chunk);
    graphar::io::write_chunk(
        &path,
        std::slice::from_ref(&batch),
        &graphar::FileType::Parquet,
    )?;
    println!("wrote {rows} rows → {}", path.display());
    Ok(())
}

/// Run one round, or poll until **Ctrl-C** — finishing the current round
/// cleanly before exiting (graceful shutdown). The library's tracing events
/// report each round (including the cumulative `metrics` event). Generic over
/// the sink so `--dry-run` and the real sink share one loop.
///
/// Between rounds the sleep is the base `interval`, but a fully-failed
/// (no-progress) round backs the next sleep off exponentially up to
/// `max_backoff` (default 10× `interval`), resetting on the first round that
/// makes progress. On shutdown, the cumulative metrics are optionally written
/// to `metrics_file` as Prometheus text. When a live `/metrics` server is
/// running (`metrics_server`), the cumulative metrics are published into its
/// shared snapshot after every round so each scrape sees the current counters.
async fn drive<Snk: GraphSink>(
    mut listener: Listener<FileSource, Snk>,
    once: bool,
    interval: u64,
    max_backoff: Option<u64>,
    metrics_file: Option<PathBuf>,
    metrics_server: &OptMetricsServer,
) -> Result<(), Box<dyn std::error::Error>> {
    if once {
        listener.sync_once().await?;
        publish_metrics(metrics_server, &listener);
        dump_metrics(&listener, &metrics_file)?;
        return Ok(());
    }
    let interval = Duration::from_secs(interval);
    let max = max_backoff
        .map(Duration::from_secs)
        .unwrap_or_else(|| interval.saturating_mul(10));
    let mut backoff = Backoff::new(interval, max);
    loop {
        let report = listener.sync_once().await?;
        // Publish the round's cumulative metrics so a live scrape sees them.
        publish_metrics(metrics_server, &listener);
        let sleep = backoff.next_sleep(report.is_failure());
        if sleep > backoff.interval() {
            tracing::warn!(
                target: "bifrost::round",
                event = "backoff",
                sleep_secs = sleep.as_secs(),
                interval_secs = backoff.interval().as_secs(),
                "round made no progress — backing off before retry",
            );
        }
        // Sleep until the next poll OR a shutdown signal — Ctrl-C during the
        // wait exits immediately, but only between rounds (the current round
        // always completes and its state is persisted first).
        tokio::select! {
            _ = tokio::signal::ctrl_c() => {
                tracing::info!(target: "bifrost", event = "shutdown", "graceful shutdown");
                dump_metrics(&listener, &metrics_file)?;
                return Ok(());
            }
            _ = tokio::time::sleep(sleep) => {}
        }
    }
}

/// Write the listener's cumulative metrics as Prometheus text to `path` (if
/// any). A `/metrics` HTTP endpoint can serve the same `render_prometheus`
/// output later without an HTTP dependency in the library.
fn dump_metrics<Snk: GraphSink>(
    listener: &Listener<FileSource, Snk>,
    path: &Option<PathBuf>,
) -> Result<(), Box<dyn std::error::Error>> {
    if let Some(path) = path {
        std::fs::write(path, knut_bifrost::render_prometheus(listener.metrics()))?;
        tracing::info!(
            target: "metrics",
            event = "metrics_written",
            path = %path.display(),
            "wrote Prometheus metrics",
        );
    }
    Ok(())
}

/// A graph sink that writes nothing — it just counts what *would* be written,
/// so `--dry-run` can exercise detect→transform with no database. The library's
/// `sink_applied` event still fires via the listener's per-table log.
#[derive(Default)]
struct NullSink;

#[async_trait]
impl GraphSink for NullSink {
    async fn apply(&mut self, mutation: &GraphMutation) -> BifrostResult<SinkReport> {
        let vertices: usize = mutation.vertices.iter().map(|v| v.batch.num_rows()).sum();
        let edges: usize = mutation.edges.iter().map(|e| e.batch.num_rows()).sum();
        let deleted: usize = mutation.deletes.iter().map(|d| d.batch.num_rows()).sum();
        tracing::info!(
            target: "bifrost::sink",
            event = "dry_run_mutation",
            vertices,
            edges,
            deleted,
            "dry-run: counted graph mutation (not written)",
        );
        Ok(SinkReport {
            vertices_written: vertices,
            edges_written: edges,
            deleted,
        })
    }
}