rbt-datalake 0.5.0

use anyhow::{bail, Context, Result};
use clap::{Parser, Subcommand, ValueEnum};
use rbt::{
    model_has_test_contract, BronzeCheckMode, OutputFormat, RbtProjectConfig, SelectMode,
    TransformationEngine,
};

use std::path::PathBuf;
use std::str::FromStr;
use std::time::Instant;
use tracing_subscriber::EnvFilter;

#[derive(Parser)]
#[command(
    name = "rbt",
    version,
    about = "Rust lake build tool: medallion SQL DAGs on Parquet/Iceberg-style tables"
)]
struct Cli {
    #[command(subcommand)]
    command: Commands,
}

#[derive(ValueEnum, Clone, Debug, PartialEq, Eq)]
enum CliFormat {
    /// Single Parquet file per model (default full-rewrite path)
    Parquet,
    Jsonl,
    Csv,
    /// Filesystem Iceberg-style table directory (data/ + metadata/)
    Iceberg,
    /// Dual-write: flat .parquet + sibling .iceberg/ table dir
    ParquetAndIceberg,
    /// Currently materializes Parquet (clone semantics reserved)
    ZeroCopyClone,
}

impl From<CliFormat> for OutputFormat {
    fn from(fmt: CliFormat) -> Self {
        match fmt {
            CliFormat::Parquet => OutputFormat::Parquet,
            CliFormat::Jsonl => OutputFormat::Jsonl,
            CliFormat::Csv => OutputFormat::Csv,
            CliFormat::Iceberg => OutputFormat::Iceberg,
            CliFormat::ParquetAndIceberg => OutputFormat::ParquetAndIceberg,
            CliFormat::ZeroCopyClone => OutputFormat::ZeroCopyClone,
        }
    }
}

#[derive(ValueEnum, Clone, Debug, PartialEq, Eq)]
enum CliBronzeCheck {
    Off,
    Warn,
    Fail,
}

impl From<CliBronzeCheck> for BronzeCheckMode {
    fn from(v: CliBronzeCheck) -> Self {
        match v {
            CliBronzeCheck::Off => BronzeCheckMode::Off,
            CliBronzeCheck::Warn => BronzeCheckMode::Warn,
            CliBronzeCheck::Fail => BronzeCheckMode::Fail,
        }
    }
}

impl FromStr for CliBronzeCheck {
    type Err = String;
    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s.to_ascii_lowercase().as_str() {
            "off" => Ok(Self::Off),
            "warn" | "warning" => Ok(Self::Warn),
            "fail" | "error" => Ok(Self::Fail),
            other => Err(format!(
                "invalid bronze-check '{}'; expected off|warn|fail",
                other
            )),
        }
    }
}

#[derive(Subcommand)]
enum Commands {
    /// Compile SQL models, build DAG, validate bronze frontmatter scan paths
    Compile {
        #[arg(short, long, default_value = ".")]
        project_dir: PathBuf,
        /// Model selector: name, +name, name+, comma-separated (dbt-like)
        #[arg(short = 's', long)]
        select: Option<String>,
        /// How to treat missing/invalid bronze scan_path: off | warn | fail
        #[arg(long, value_enum, default_value_t = CliBronzeCheck::Warn)]
        bronze_check: CliBronzeCheck,
    },
    /// Execute transformation pipeline DAG with frontmatter-driven bronze registration
    Run {
        #[arg(short, long, default_value = ".")]
        project_dir: PathBuf,
        /// Model selector: name, +name, name+, comma-separated. Ancestors always included.
        #[arg(short = 's', long)]
        select: Option<String>,
        #[arg(short, long, default_value = "./target/output")]
        output_dir: PathBuf,
        /// Output format override for all models in this run
        #[arg(short, long, value_enum, default_value_t = CliFormat::Parquet)]
        format: CliFormat,
        /// Pre-flight bronze path check before execution (default: fail)
        #[arg(long, value_enum, default_value_t = CliBronzeCheck::Fail)]
        bronze_check: CliBronzeCheck,
    },
    /// Run frontmatter tests for selected models (executes subgraph then asserts)
    Test {
        #[arg(short, long, default_value = ".")]
        project_dir: PathBuf,
        /// Model selector (default: all models that declare tests or grain/unique_key)
        #[arg(short = 's', long)]
        select: Option<String>,
        /// Optional single model (alias for --select)
        #[arg(short, long)]
        model: Option<String>,
        #[arg(short, long, default_value = "./target/output")]
        output_dir: PathBuf,
        #[arg(short, long, value_enum, default_value_t = CliFormat::Parquet)]
        format: CliFormat,
        #[arg(long, value_enum, default_value_t = CliBronzeCheck::Fail)]
        bronze_check: CliBronzeCheck,
    },
    /// Static validation: load project, build DAG, bronze paths, layer rules (no execute)
    Validate {
        #[arg(short, long, default_value = ".")]
        project_dir: PathBuf,
        #[arg(short = 's', long)]
        select: Option<String>,
        #[arg(long, value_enum, default_value_t = CliBronzeCheck::Fail)]
        bronze_check: CliBronzeCheck,
        /// Emit JSON report to stdout (machine-readable)
        #[arg(long, default_value_t = false)]
        json: bool,
    },
    /// Explain a model: compiled SQL, deps, layer, bronze contract, output path
    Explain {
        #[arg(short, long, default_value = ".")]
        project_dir: PathBuf,
        /// Model name (required)
        #[arg(short = 's', long)]
        select: String,
        #[arg(long, default_value_t = false)]
        json: bool,
    },
    /// Preview model rows: materialize ancestors, run target SQL with LIMIT (no target write)
    Preview {
        #[arg(short, long, default_value = ".")]
        project_dir: PathBuf,
        /// Single model name
        #[arg(short = 's', long)]
        select: String,
        #[arg(short, long, default_value = "./target/output")]
        output_dir: PathBuf,
        #[arg(long, default_value_t = 20)]
        limit: usize,
        #[arg(short, long, value_enum, default_value_t = CliFormat::Parquet)]
        format: CliFormat,
        #[arg(long, value_enum, default_value_t = CliBronzeCheck::Fail)]
        bronze_check: CliBronzeCheck,
    },
    /// Thesis measure packs (wall time, rows, optional RSS) — honest experiment harness
    Measure {
        #[arg(short, long, default_value = ".")]
        project_dir: PathBuf,
        /// Scenario: smoke_pipeline | validate_dx | incremental_append
        #[arg(long, default_value = "smoke_pipeline")]
        scenario: String,
        #[arg(short, long, default_value = "./target/output")]
        output_dir: PathBuf,
        /// Write JSON report (default: {project}/.rbt/measure/{scenario}.json)
        #[arg(long)]
        report: Option<PathBuf>,
        /// Print JSON report to stdout
        #[arg(long, default_value_t = false)]
        json: bool,
    },
    /// Run micro-benchmarks measuring transformation engine throughput (rows/sec)
    Bench {
        #[arg(short, long, default_value_t = 1000000)]
        num_rows: usize,
    },
}

#[tokio::main]
async fn main() -> Result<()> {
    tracing_subscriber::fmt()
        .with_env_filter(
            EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("info")),
        )
        .init();

    let cli = Cli::parse();

    match cli.command {
        Commands::Compile {
            project_dir,
            select,
            bronze_check,
        } => {
            println!(
                "[rbt] Compiling project {:?} (select={:?}, bronze_check={})...",
                project_dir,
                select,
                BronzeCheckMode::from(bronze_check.clone())
            );
            let config = RbtProjectConfig::load(&project_dir)?;
            let full = config.build_dag(&project_dir, None)?;
            let dag = full.apply_select(select.as_deref(), SelectMode::Exact)?;
            let tiers = dag.execution_tiers()?;
            println!(
                "[rbt] DAG built with {} model(s) in {} tier(s):",
                dag.node_map.len(),
                tiers.len()
            );
            for (i, tier) in tiers.iter().enumerate() {
                let names: Vec<&str> = tier.iter().map(|m| m.name.as_str()).collect();
                println!("  Tier {}: {:?}", i, names);
            }

            // Bronze validation on full project (or selected subgraph for paths)
            let report = dag.validate_bronze_sources_with_roots(
                &project_dir,
                bronze_check.into(),
                &config.roots,
            )?;
            for d in &report.diagnostics {
                eprintln!("{d}");
            }
            if report.has_errors() {
                bail!(
                    "[rbt] compile failed: {} bronze error(s), {} warning(s)",
                    report.error_count(),
                    report.warning_count()
                );
            }
            if report.warning_count() > 0 {
                println!(
                    "[rbt] compile succeeded with {} bronze warning(s)",
                    report.warning_count()
                );
            } else {
                println!("[rbt] compile succeeded (bronze sources ok)");
            }
        }
        Commands::Run {
            project_dir,
            select,
            output_dir,
            format,
            bronze_check,
        } => {
            println!(
                "[rbt] Executing pipeline from {:?} (select: {:?}, format: {:?}, output: {:?}, bronze_check={:?})...",
                project_dir, select, format, output_dir, bronze_check
            );
            let start = Instant::now();

            let config = RbtProjectConfig::load(&project_dir)?;
            let full = config.build_dag(&project_dir, Some(format.into()))?;
            let dag = full
                .apply_select(select.as_deref(), SelectMode::Execute)
                .context("invalid --select")?;

            let report = dag.validate_bronze_sources_with_roots(
                &project_dir,
                bronze_check.into(),
                &config.roots,
            )?;
            for d in &report.diagnostics {
                eprintln!("{d}");
            }
            if report.has_errors() {
                bail!(
                    "[rbt] run aborted: {} bronze error(s) — fix scan_path/frontmatter or use --bronze-check=warn",
                    report.error_count()
                );
            }

            println!(
                "[rbt] Running {} model(s): {:?}",
                dag.node_map.len(),
                dag.topological_sequence()?
                    .iter()
                    .map(|m| m.name.as_str())
                    .collect::<Vec<_>>()
            );

            let engine = TransformationEngine::new();
            let summary = engine.execute_dag(&dag, &project_dir, &output_dir).await?;
            let duration = start.elapsed();

            println!(
                "[rbt] Completed {} models ({} rows, {} bronze sources) in {:.2?}",
                summary.models_executed,
                summary.total_rows_produced,
                summary.bronze_sources_registered,
                duration
            );
        }
        Commands::Test {
            project_dir,
            select,
            model,
            output_dir,
            format,
            bronze_check,
        } => {
            let select = match (select, model) {
                (Some(s), Some(m)) => Some(format!("{},{}", s, m)),
                (Some(s), None) => Some(s),
                (None, Some(m)) => Some(m),
                (None, None) => None,
            };

            println!(
                "[rbt] Testing project {:?} (select={:?})...",
                project_dir, select
            );

            let config = RbtProjectConfig::load(&project_dir)?;
            let full = config.build_dag(&project_dir, Some(format.into()))?;

            // Default: models that declare tests or grain/unique_key
            let select_spec = if let Some(s) = select {
                s
            } else {
                let names = full.models_with_test_contract()?;
                if names.is_empty() {
                    println!("[rbt] No models with frontmatter tests/grain/unique_key found.");
                    return Ok(());
                }
                names.join(",")
            };

            let dag = full
                .apply_select(Some(&select_spec), SelectMode::Execute)
                .context("invalid --select / --model")?;

            let report = dag.validate_bronze_sources_with_roots(
                &project_dir,
                bronze_check.into(),
                &config.roots,
            )?;
            for d in &report.diagnostics {
                eprintln!("{d}");
            }
            if report.has_errors() {
                bail!(
                    "[rbt] test aborted: {} bronze error(s)",
                    report.error_count()
                );
            }

            println!(
                "[rbt] Executing {} model(s) then validating frontmatter tests...",
                dag.node_map.len()
            );

            let engine = TransformationEngine::new();
            // execute_dag already runs frontmatter tests and fails hard on error
            let summary = engine.execute_dag(&dag, &project_dir, &output_dir).await?;

            // Summarize which models had tests (already enforced during materialize)
            let mut tested = 0usize;
            let mut without = 0usize;
            for node in dag.topological_sequence()? {
                if model_has_test_contract(&node) {
                    tested += 1;
                    println!(
                        "  PASS  {} (frontmatter tests executed during materialize)",
                        node.name
                    );
                } else {
                    without += 1;
                    println!("  SKIP  {} (no tests/grain/unique_key)", node.name);
                }
            }

            println!(
                "[rbt] Tests finished: {} models with tests, {} skipped, {} rows produced across run",
                tested, without, summary.total_rows_produced
            );
        }
        Commands::Validate {
            project_dir,
            select,
            bronze_check,
            json,
        } => {
            let config = RbtProjectConfig::load(&project_dir)?;
            let full = config.build_dag(&project_dir, None)?;
            let dag = full.apply_select(select.as_deref(), SelectMode::Exact)?;
            let tiers = dag.execution_tiers()?;
            let report = dag.validate_bronze_sources_with_roots(
                &project_dir,
                bronze_check.into(),
                &config.roots,
            )?;

            let mut issues: Vec<String> = Vec::new();
            for d in &report.diagnostics {
                issues.push(d.to_string());
            }
            // Layer / ref hygiene: every ref must resolve in full project
            for node in full.topological_sequence()? {
                for dep in &node.dependencies {
                    if let rbt::DependencyRef::Model(name) = dep {
                        if !full.node_map.contains_key(name) {
                            issues.push(format!(
                                "E_RBT_VALIDATE_REF: model '{}' refs unknown model '{}'",
                                node.name, name
                            ));
                        }
                    }
                }
            }

            let ok = !report.has_errors()
                && !issues.iter().any(|i| i.contains("E_RBT_VALIDATE_REF"));

            if json {
                let body = serde_json::json!({
                    "ok": ok,
                    "project": config.name,
                    "models": dag.node_map.len(),
                    "tiers": tiers.len(),
                    "bronze_errors": report.error_count(),
                    "bronze_warnings": report.warning_count(),
                    "issues": issues,
                    "tier_plan": tiers.iter().map(|t| {
                        t.iter().map(|m| m.name.clone()).collect::<Vec<_>>()
                    }).collect::<Vec<_>>(),
                });
                println!("{}", serde_json::to_string_pretty(&body)?);
            } else {
                println!(
                    "[rbt] validate project {:?} ({} model(s), {} tier(s))",
                    project_dir,
                    dag.node_map.len(),
                    tiers.len()
                );
                for (i, tier) in tiers.iter().enumerate() {
                    let names: Vec<&str> = tier.iter().map(|m| m.name.as_str()).collect();
                    println!("  Tier {i}: {names:?}");
                }
                for d in &report.diagnostics {
                    eprintln!("{d}");
                }
                for i in &issues {
                    if i.starts_with("E_RBT_VALIDATE") {
                        eprintln!("{i}");
                    }
                }
                if ok {
                    println!(
                        "[rbt] validate OK (bronze errors={}, warnings={})",
                        report.error_count(),
                        report.warning_count()
                    );
                } else {
                    bail!(
                        "[rbt] validate FAILED (bronze errors={}, issues={})",
                        report.error_count(),
                        issues.len()
                    );
                }
            }
            if !ok && json {
                std::process::exit(1);
            }
        }
        Commands::Explain {
            project_dir,
            select,
            json,
        } => {
            let config = RbtProjectConfig::load(&project_dir)?;
            let full = config.build_dag(&project_dir, None)?;
            let name = select.trim();
            let node = full
                .topological_sequence()?
                .into_iter()
                .find(|m| m.name == name)
                .ok_or_else(|| anyhow::anyhow!("E_RBT_EXPLAIN: unknown model '{name}'"))?;

            let deps: Vec<String> = node
                .dependencies
                .iter()
                .map(|d| match d {
                    rbt::DependencyRef::Model(n) => format!("ref:{n}"),
                    rbt::DependencyRef::Source {
                        source_name,
                        table_name,
                    } => format!("source:{source_name}.{table_name}"),
                })
                .collect();
            let fm = node.frontmatter.as_ref();
            if json {
                let body = serde_json::json!({
                    "name": node.name,
                    "layer": format!("{:?}", node.layer),
                    "materialization": format!("{:?}", node.materialization),
                    "output_format": format!("{:?}", node.output_format),
                    "output_path": node.output_path,
                    "dependencies": deps,
                    "description": node.description,
                    "compiled_sql": node.compiled_sql,
                    "bronze": fm.map(|f| serde_json::json!({
                        "scan_path": f.scan_path,
                        "source_format": f.source_format.as_ref().map(|s| s.as_str()),
                        "path_glob": f.path_glob,
                        "partition_by": f.partition_by,
                        "require_partitions": f.require_partitions,
                    })),
                });
                println!("{}", serde_json::to_string_pretty(&body)?);
            } else {
                println!("[rbt] explain model '{}'", node.name);
                println!("  layer:           {:?}", node.layer);
                println!("  materialization: {:?}", node.materialization);
                println!("  output_format:   {:?}", node.output_format);
                if let Some(p) = &node.output_path {
                    println!("  output_path:     {p}");
                }
                println!("  dependencies:    {deps:?}");
                if let Some(f) = fm {
                    if let Some(sp) = &f.scan_path {
                        println!("  bronze.scan_path: {sp}");
                    }
                    if let Some(sf) = &f.source_format {
                        println!("  bronze.format:    {}", sf.as_str());
                    }
                    if let Some(g) = &f.path_glob {
                        println!("  bronze.path_glob: {g:?}");
                    }
                }
                println!("--- compiled SQL ---");
                println!("{}", node.compiled_sql);
            }
        }
        Commands::Preview {
            project_dir,
            select,
            output_dir,
            limit,
            format,
            bronze_check,
        } => {
            let config = RbtProjectConfig::load(&project_dir)?;
            let full = config.build_dag(&project_dir, Some(format.into()))?;
            let report = full.validate_bronze_sources_with_roots(
                &project_dir,
                bronze_check.into(),
                &config.roots,
            )?;
            for d in &report.diagnostics {
                eprintln!("{d}");
            }
            if report.has_errors() {
                bail!(
                    "[rbt] preview aborted: {} bronze error(s)",
                    report.error_count()
                );
            }
            println!(
                "[rbt] preview model '{}' (limit={limit}, ancestors materialize if needed)...",
                select
            );
            let engine = TransformationEngine::new();
            let result = engine
                .preview_model(&full, &project_dir, &output_dir, select.trim(), limit)
                .await?;
            println!(
                "[rbt] preview '{}': {} row(s) (limit {}), {} ancestor model(s) executed",
                result.model, result.rows, result.limit, result.ancestors_executed
            );
            // Print a simple table: schema + first values
            if let Some(batch) = result.batches.first() {
                let schema = batch.schema();
                let cols: Vec<&str> = schema.fields().iter().map(|f| f.name().as_str()).collect();
                println!("columns: {cols:?}");
                let n = batch.num_rows().min(limit);
                for row in 0..n {
                    let mut cells = Vec::new();
                    for c in 0..batch.num_columns() {
                        let arr = batch.column(c);
                        cells.push(arrow_cell_display(arr, row));
                    }
                    println!("  row[{row}]: {}", cells.join(" | "));
                }
                if result.batches.len() > 1 {
                    println!(
                        "  … {} additional batch(es) not printed",
                        result.batches.len() - 1
                    );
                }
            } else {
                println!("  (no rows)");
            }
        }
        Commands::Measure {
            project_dir,
            scenario,
            output_dir,
            report,
            json,
        } => {
            println!(
                "[rbt] measure scenario='{scenario}' project={project_dir:?}"
            );
            let report_data =
                rbt::run_measure_scenario(&scenario, &project_dir, &output_dir).await?;
            let out = report.unwrap_or_else(|| {
                rbt::default_report_path(&project_dir, &scenario)
            });
            rbt::write_measure_report(&report_data, &out)?;
            if json {
                println!("{}", serde_json::to_string_pretty(&report_data)?);
            } else {
                println!(
                    "[rbt] measure OK={} wall_ms={} models={} rows={} bronze={} rss_kb={:?}",
                    report_data.ok,
                    report_data.wall_ms,
                    report_data.models_executed,
                    report_data.total_rows,
                    report_data.bronze_sources,
                    report_data.peak_rss_kb
                );
                println!("[rbt] report written → {}", out.display());
                for n in &report_data.notes {
                    println!("  note: {n}");
                }
            }
            if !report_data.ok {
                bail!(
                    "[rbt] measure scenario failed: {}",
                    report_data.error.unwrap_or_default()
                );
            }
        }
        Commands::Bench { num_rows } => {
            println!(
                "[rbt BENCHMARK] Generating & transforming {} rows in-memory...",
                num_rows
            );
            let start = Instant::now();

            let engine = TransformationEngine::new();
            let query = format!(
                "SELECT id, id * 2 AS val, CONCAT('user_', CAST(id AS VARCHAR)) AS name FROM (SELECT \"range()\".value AS id FROM range(0, {}))",
                num_rows
            );

            let df = engine.ctx.sql(&query).await?;
            let batches = df.collect().await?;
            let duration = start.elapsed();

            let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
            let rows_per_sec = (total_rows as f64) / duration.as_secs_f64();

            println!(
                "[rbt BENCHMARK RESULT] Processed {} rows in {:.2?} ({:.0} rows/sec)",
                total_rows, duration, rows_per_sec
            );
        }
    }

    Ok(())
}

fn arrow_cell_display(array: &arrow::array::ArrayRef, row: usize) -> String {
    use arrow::array::Array;
    use arrow::util::display::{ArrayFormatter, FormatOptions};
    if array.is_null(row) {
        return "NULL".into();
    }
    let opts = FormatOptions::default().with_display_error(true);
    if let Ok(fmt) = ArrayFormatter::try_new(array.as_ref(), &opts) {
        return fmt.value(row).to_string();
    }
    "?".to_string()
}