xlog/

main.rs

use clap::{Parser, Subcommand, ValueEnum};
use std::collections::HashMap;
use std::io::Read;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::time::Duration;

use arrow::csv::WriterBuilder;
use arrow::util::pretty::pretty_format_batches;
use xlog_core::{symbol, MemoryBudget, Result, XlogError};
use xlog_cuda::{CudaDevice, CudaKernelProvider, GpuMemoryManager};
use xlog_gpu::logic::{normalize_program_for_execution, LogicProgram};
use xlog_ir::{EirBodyLiteral, EirTerm};
use xlog_logic::ast::{BodyLiteral, ProbEngine, Program};
use xlog_logic::compile::load_modules;
#[cfg(feature = "host-io")]
use xlog_logic::parse_program;
use xlog_logic::IncrementalParseResult;
use xlog_logic::{rewrite_magic_sets, MagicSetReport, MagicSetStatus, ParserSession};
use xlog_logic::{stratify, Compiler};
use xlog_logic::{QueryProofTrace, RuleProvenance};
#[cfg(feature = "host-io")]
use xlog_prob::exact::ExactDdnnfProgram;
#[cfg(feature = "host-io")]
use xlog_prob::exact::GpuConfig;
#[cfg(feature = "host-io")]
use xlog_prob::mc::{McEvalConfig, McProgram, McSamplingMethod};
use xlog_prob::provenance::{AggregateLiftReport, Value};

mod generated_rule_diagnostics;

use generated_rule_diagnostics::{
    explain_generated_rule_diagnostics, print_generated_rule_diagnostics_json,
    GeneratedRuleDiagnostic,
};

#[derive(Parser)]
#[command(author, version, about = "XLOG CLI")]
pub struct Cli {
    #[command(subcommand)]
    command: Command,
}

#[derive(Subcommand)]
enum Command {
    Run(RunArgs),
    Prob(ProbArgs),
    Explain(ExplainArgs),
    Repl(ReplArgs),
    Watch(WatchArgs),
}

#[derive(Parser)]
struct RunArgs {
    source: PathBuf,
    #[arg(long, default_value = "0")]
    device: usize,
    #[arg(long, default_value = "1024")]
    memory_mb: u64,
    #[arg(long)]
    input: Vec<String>,
    #[arg(long, value_enum, default_value = "pretty")]
    output: OutputFormat,
    #[arg(long)]
    output_dir: Option<PathBuf>,
    /// Show execution statistics (timing, memory usage)
    #[arg(long)]
    stats: bool,
    /// Stats output format (human or json)
    #[arg(long, value_enum, default_value = "human")]
    stats_format: StatsFormat,
    /// Additional directories to search for modules (colon-separated)
    #[arg(long, value_delimiter = ':')]
    module_path: Vec<PathBuf>,
    /// Dump the compiled epistemic execution plan (EIR-derived GPU plan, world-view
    /// integrity constraints, and fail-closed execution policy) as JSON to this path.
    /// No-op for ordinary (non-epistemic) programs. This compiled
    /// epistemic-plan/EIR JSON dump exposes accepted `know`/`possible` literals
    /// and lets a caller verify that unsupported execution shapes are rejected.
    #[arg(long)]
    epistemic_plan_json: Option<PathBuf>,
    /// Engage the worst-case-optimal join (WCOJ) subsystem for eligible
    /// multiway rules (triangle + 4-cycle). Without this the deterministic
    /// runner uses binary joins, which blow up on skewed cyclic queries.
    /// Sets the documented WCOJ dispatch gates for this process.
    #[arg(long)]
    wcoj: bool,
}

#[derive(Copy, Clone, ValueEnum, Default)]
enum StatsFormat {
    #[default]
    Human,
    Json,
}

#[derive(Parser)]
struct ProbArgs {
    source: PathBuf,
    #[arg(long, default_value = "0")]
    device: usize,
    #[arg(long, default_value = "1024")]
    memory_mb: u64,
    #[arg(long, value_enum)]
    prob_engine: Option<ProbEngineCli>,
    #[arg(long)]
    samples: Option<usize>,
    #[arg(long)]
    seed: Option<u64>,
    #[arg(long)]
    confidence: Option<f64>,
    #[arg(long, value_enum)]
    prob_method: Option<ProbMethodCli>,
    #[arg(long, alias = "max-nonmonotone-iterations")]
    prob_max_nonmonotone_iterations: Option<usize>,
    /// Allow the labeled CPU oracle when the resident GPU MC engine rejects
    /// the program (negation, aggregates, ...). Fail-closed when unset; the
    /// result is labeled `mc_engine: cpu-oracle` and is not GPU-native evidence.
    #[arg(long)]
    allow_cpu_oracle: bool,
    #[arg(long, value_enum, default_value = "pretty")]
    output: ProbOutputFormat,
    #[arg(long)]
    output_dir: Option<PathBuf>,
    /// Additional directories to search for modules (colon-separated)
    #[arg(long, value_delimiter = ':')]
    module_path: Vec<PathBuf>,
}

#[derive(Parser)]
struct ExplainArgs {
    source: PathBuf,
    #[arg(long, value_enum, default_value = "text")]
    format: ExplainFormat,
    /// Additional directories to search for modules (colon-separated)
    #[arg(long, value_delimiter = ':')]
    module_path: Vec<PathBuf>,
}

#[derive(Parser)]
struct ReplArgs {
    /// Additional directories to search for modules (colon-separated)
    #[arg(long, value_delimiter = ':')]
    module_path: Vec<PathBuf>,
}

#[derive(Parser)]
struct WatchArgs {
    source: PathBuf,
    #[arg(long, default_value = "250")]
    debounce_ms: u64,
    #[arg(long)]
    explain: bool,
    #[arg(long)]
    once: bool,
    /// Additional directories to search for modules (colon-separated)
    #[arg(long, value_delimiter = ':')]
    module_path: Vec<PathBuf>,
}

#[derive(Copy, Clone, ValueEnum)]
enum ExplainFormat {
    Text,
    Json,
    Dot,
}

#[derive(Copy, Clone, ValueEnum)]
enum OutputFormat {
    Pretty,
    Csv,
    Arrow,
}

#[derive(Copy, Clone, ValueEnum)]
enum ProbOutputFormat {
    Pretty,
    Csv,
    Arrow,
    Json,
}

#[derive(Copy, Clone, ValueEnum)]
enum ProbEngineCli {
    #[value(name = "exact_ddnnf")]
    ExactDdnnf,
    Mc,
}

#[derive(Copy, Clone, ValueEnum)]
enum ProbMethodCli {
    Rejection,
    #[value(name = "evidence_clamping")]
    EvidenceClamping,
}

fn main() -> Result<()> {
    let cli = Cli::parse();
    match cli.command {
        Command::Run(args) => run_deterministic(args),
        Command::Prob(args) => run_probabilistic(args),
        Command::Explain(args) => explain(args),
        Command::Repl(args) => repl(args),
        Command::Watch(args) => watch(args),
    }
}

fn explain(args: ExplainArgs) -> Result<()> {
    let source = std::fs::read_to_string(&args.source).map_err(|e| {
        XlogError::Execution(format!("Failed to read {}: {}", args.source.display(), e))
    })?;
    let mut parser_session = ParserSession::new();
    let parsed = parser_session.parse_path(&args.source, &source)?;
    let parsed = resolve_explain_imports(parsed, &args.source, args.module_path)?;
    let report = build_explain_report(parsed, Some(&args.source))?;
    match args.format {
        ExplainFormat::Text => print_explain_text(&report),
        ExplainFormat::Json => print_explain_json(&report),
        ExplainFormat::Dot => print_magic_dot(&report.magic_sets),
    }
    Ok(())
}

fn resolve_explain_imports(
    mut parsed: IncrementalParseResult,
    source_path: &Path,
    module_path: Vec<PathBuf>,
) -> Result<IncrementalParseResult> {
    parsed.program = resolve_program_imports(
        parsed.program,
        source_path,
        module_path,
        ModuleMergeErrorKind::Execution,
    )?;
    Ok(parsed)
}

#[derive(Clone, Copy)]
enum ModuleMergeErrorKind {
    Execution,
    Compilation,
}

fn resolve_program_imports(
    program: Program,
    source_path: &Path,
    module_path: Vec<PathBuf>,
    merge_error_kind: ModuleMergeErrorKind,
) -> Result<Program> {
    if program.imports.is_empty() {
        return Ok(program);
    }
    let resolver = load_modules(source_path, module_path)
        .map_err(|e| XlogError::Execution(format!("Module resolution failed: {}", e)))?;
    warn_ignored_import_pragmas(&resolver);
    resolver
        .merge_imports(program)
        .map_err(|error| match merge_error_kind {
            ModuleMergeErrorKind::Execution => {
                XlogError::Execution(format!("Module resolution failed: {error}"))
            }
            ModuleMergeErrorKind::Compilation => {
                XlogError::Compilation(format!("Module resolution failed: {error}"))
            }
        })
}

/// Surface pragmas declared in imported modules on stderr. Pragmas are
/// entry-file-scoped, so these directives are dropped at merge time; the
/// warning keeps that scoping from being silent.
fn warn_ignored_import_pragmas(resolver: &xlog_logic::resolver::ModuleResolver) {
    for warning in resolver.ignored_import_pragmas() {
        eprintln!("{}", warning);
    }
}

fn repl(args: ReplArgs) -> Result<()> {
    let _ = args.module_path;
    let mut input = String::new();
    std::io::stdin()
        .read_to_string(&mut input)
        .map_err(|e| XlogError::Execution(format!("Failed to read stdin: {}", e)))?;
    let mut parser_session = ParserSession::new();
    let parsed = parser_session.parse_path("<repl>", &input)?;
    println!(
        "repl: statements={} cache_hits={} cache_misses={}",
        parsed.stats.statement_count, parsed.stats.hits, parsed.stats.misses
    );
    println!(
        "state: rules={} queries={} prob_queries={}",
        parsed.program.rules.len(),
        parsed.program.queries.len(),
        parsed.program.prob_queries.len()
    );
    Ok(())
}

fn watch(args: WatchArgs) -> Result<()> {
    let mut parser_session = ParserSession::new();
    loop {
        let source = std::fs::read_to_string(&args.source).map_err(|e| {
            XlogError::Execution(format!("Failed to read {}: {}", args.source.display(), e))
        })?;
        let parsed = parser_session.parse_path(&args.source, &source)?;
        let parsed = if args.explain {
            resolve_explain_imports(parsed, &args.source, args.module_path.clone())?
        } else {
            parsed
        };
        println!(
            "watch: statements={} cache_hits={} cache_misses={}",
            parsed.stats.statement_count, parsed.stats.hits, parsed.stats.misses
        );
        if args.explain {
            let report = build_explain_report(parsed, Some(&args.source))?;
            print_explain_text(&report);
        }
        if args.once {
            break;
        }
        std::thread::sleep(Duration::from_millis(args.debounce_ms));
    }
    Ok(())
}

struct ExplainReport {
    program: Program,
    parse_stats: xlog_logic::ParseCacheStats,
    epistemic: serde_json::Value,
    magic_sets: MagicSetReport,
    aggregate_lifting: Vec<AggregateLiftReport>,
    generated_rule_diagnostics: Vec<GeneratedRuleDiagnostic>,
    generated_rule_diagnostics_status: String,
    generated_rule_diagnostics_reason: Option<String>,
    rule_provenance: Vec<RuleProvenance>,
    proof_traces: Vec<QueryProofTrace>,
    stratification_status: String,
    stratification_reason: Option<String>,
    stratification_count: usize,
    aggregate_lifting_status: String,
    aggregate_lifting_reason: Option<String>,
    rir_status: String,
    rir_reason: Option<String>,
    rir_sccs: usize,
    optimizer_status: String,
    optimizer_reason: Option<String>,
    optimizer_memory_peak: u64,
    wcoj_status: String,
    wcoj_reason: Option<String>,
}

fn build_explain_report(
    parsed: xlog_logic::IncrementalParseResult,
    source_path: Option<&Path>,
) -> Result<ExplainReport> {
    let source_program = parsed.program;
    match normalize_program_for_execution(source_program.clone()) {
        Ok(analysis_program) => {
            let magic_rewrite = rewrite_magic_sets(&analysis_program)?;
            let (rule_provenance, proof_traces) = explain_source_diagnostics(
                &source_program,
                &analysis_program,
                &magic_rewrite.program,
            );
            let aggregate_lifting = explain_aggregate_lifting(&analysis_program)?;
            let epistemic = explain_epistemic(&analysis_program);
            let (stratification_status, stratification_count) = match stratify(&analysis_program) {
                Ok(strata) => ("ok".to_string(), strata.len()),
                Err(error) => (format!("error: {error}"), 0),
            };
            let mut compiler = Compiler::new();
            let (
                rir_status,
                rir_reason,
                rir_sccs,
                optimizer_status,
                optimizer_reason,
                optimizer_memory_peak,
            ) = match compiler.compile_program(&analysis_program) {
                Ok(plan) => (
                    "ok".to_string(),
                    None,
                    plan.sccs.len(),
                    "ok".to_string(),
                    None,
                    plan.est_memory_peak,
                ),
                Err(error) => {
                    let reason = format!("RIR compilation failed: {error}");
                    (
                        format!("error: {error}"),
                        Some(reason.clone()),
                        0,
                        "not_available".to_string(),
                        Some(reason),
                        0,
                    )
                }
            };
            let (
                generated_rule_diagnostics,
                generated_rule_diagnostics_status,
                generated_rule_diagnostics_reason,
            ) = if let Some(reason) = &rir_reason {
                (
                    Vec::new(),
                    "not_available".to_string(),
                    Some(reason.clone()),
                )
            } else {
                match explain_generated_rule_diagnostics(
                    &source_program,
                    &analysis_program,
                    source_path,
                ) {
                    Ok(diagnostics) => (diagnostics, "ok".to_string(), None),
                    Err(error) => {
                        let reason = format!("generated-rule diagnostics failed: {error}");
                        (Vec::new(), "not_available".to_string(), Some(reason))
                    }
                }
            };
            let (wcoj_status, wcoj_reason) = if let Some(reason) = &rir_reason {
                ("not_available".to_string(), Some(reason.clone()))
            } else {
                ("reported".to_string(), None)
            };
            Ok(ExplainReport {
                program: source_program,
                parse_stats: parsed.stats,
                epistemic,
                magic_sets: magic_rewrite.report,
                aggregate_lifting,
                generated_rule_diagnostics,
                generated_rule_diagnostics_status,
                generated_rule_diagnostics_reason,
                rule_provenance,
                proof_traces,
                stratification_status,
                stratification_reason: None,
                stratification_count,
                aggregate_lifting_status: "ok".to_string(),
                aggregate_lifting_reason: None,
                rir_status,
                rir_reason,
                rir_sccs,
                optimizer_status,
                optimizer_reason,
                optimizer_memory_peak,
                wcoj_status,
                wcoj_reason,
            })
        }
        Err(error) => {
            let normalization_reason = format!("execution normalization failed: {error}");
            let rule_provenance = xlog_logic::rule_provenance(&source_program, None);
            let proof_traces = xlog_logic::query_proof_traces(&source_program, &rule_provenance);
            Ok(ExplainReport {
                program: source_program,
                parse_stats: parsed.stats,
                epistemic: unavailable_epistemic_analysis(&normalization_reason),
                magic_sets: MagicSetReport {
                    status: MagicSetStatus::Declined,
                    generated_predicates: Vec::new(),
                    adorned_predicates: Vec::new(),
                    declined_reasons: vec![normalization_reason.clone()],
                },
                aggregate_lifting: Vec::new(),
                generated_rule_diagnostics: Vec::new(),
                generated_rule_diagnostics_status: "not_available".to_string(),
                generated_rule_diagnostics_reason: Some(normalization_reason.clone()),
                rule_provenance,
                proof_traces,
                stratification_status: "not_available".to_string(),
                stratification_reason: Some(normalization_reason.clone()),
                stratification_count: 0,
                aggregate_lifting_status: "not_available".to_string(),
                aggregate_lifting_reason: Some(normalization_reason.clone()),
                rir_status: "not_available".to_string(),
                rir_reason: Some(normalization_reason.clone()),
                rir_sccs: 0,
                optimizer_status: "not_available".to_string(),
                optimizer_reason: Some(normalization_reason.clone()),
                optimizer_memory_peak: 0,
                wcoj_status: "not_available".to_string(),
                wcoj_reason: Some(normalization_reason),
            })
        }
    }
}

fn unavailable_epistemic_analysis(reason: &str) -> serde_json::Value {
    let unavailable = serde_json::json!({
        "status": "not_available",
        "reason": reason,
    });
    serde_json::json!({
        "eir": unavailable.clone(),
        "gpu_plan": unavailable.clone(),
        "executable_plan": unavailable,
    })
}

fn explain_source_diagnostics(
    source_program: &Program,
    analysis_program: &Program,
    rewritten_program: &Program,
) -> (Vec<RuleProvenance>, Vec<QueryProofTrace>) {
    xlog_logic::source_diagnostics(source_program, analysis_program, Some(rewritten_program))
}

fn explain_epistemic(program: &Program) -> serde_json::Value {
    if !program_has_epistemic_literals(program) {
        let not_applicable = serde_json::json!({
            "status": "not_applicable",
            "reason": "program has no epistemic literals",
            "epistemic_literal_count": 0,
        });
        return serde_json::json!({
            "eir": not_applicable.clone(),
            "gpu_plan": not_applicable.clone(),
            "executable_plan": not_applicable,
        });
    }

    let eir = match xlog_logic::build_eir(program) {
        Ok(eir) => {
            let literals = eir
                .rules
                .iter()
                .enumerate()
                .flat_map(|(rule_index, rule)| {
                    rule.body.iter().filter_map(move |lit| match lit {
                        EirBodyLiteral::Epistemic(epistemic) => Some(serde_json::json!({
                            "rule_index": rule_index,
                            "literal": eir_epistemic_literal_json(epistemic),
                        })),
                        _ => None,
                    })
                })
                .collect::<Vec<_>>();
            let rule_summaries = eir
                .rules
                .iter()
                .enumerate()
                .map(|(rule_index, rule)| {
                    let epistemic_literal_count = rule
                        .body
                        .iter()
                        .filter(|lit| matches!(lit, EirBodyLiteral::Epistemic(_)))
                        .count();
                    let relational_body_atoms = rule
                        .body
                        .iter()
                        .filter(|lit| {
                            matches!(lit, EirBodyLiteral::Relational { negated: false, .. })
                        })
                        .count();
                    serde_json::json!({
                        "rule_index": rule_index,
                        "head": eir_atom_json(&rule.head),
                        "body_literal_count": rule.body.len(),
                        "epistemic_literal_count": epistemic_literal_count,
                        "relational_body_atoms": relational_body_atoms,
                    })
                })
                .collect::<Vec<_>>();
            serde_json::json!({
                "status": "ok",
                "mode": format!("{:?}", eir.mode),
                "rule_count": eir.rules.len(),
                "epistemic_literal_count": literals.len(),
                "literals": literals,
                "rules": rule_summaries,
            })
        }
        Err(err) => serde_json::json!({
            "status": "error",
            "error": err.to_string(),
        }),
    };

    let gpu_plan = match xlog_logic::epistemic::plan_epistemic_gpu_execution(program) {
        Ok(plan) => {
            let reductions = plan
                .reductions
                .iter()
                .map(|reduction| {
                    serde_json::json!({
                        "rule_index": reduction.rule_index,
                        "head_predicate": &reduction.head_predicate,
                        "relational_body_atoms": reduction.relational_body_atoms,
                        "wcoj_status": format!("{:?}", reduction.wcoj_status),
                    })
                })
                .collect::<Vec<_>>();
            let tuple_membership = plan
                .tuple_membership_bindings
                .iter()
                .map(|binding| {
                    serde_json::json!({
                        "literal_index": binding.literal_index,
                        "reduction_index": binding.reduction_index,
                        "predicate": &binding.predicate,
                        "arity": binding.arity,
                        "key_columns": binding.key_columns,
                        "key_terms": binding.key_terms.iter().map(eir_term_label).collect::<Vec<_>>(),
                        "bound_output_columns": binding.bound_output_columns,
                        "op": format!("{:?}", binding.op),
                        "negated": binding.negated,
                    })
                })
                .collect::<Vec<_>>();
            serde_json::json!({
                "status": "ok",
                "mode": format!("{:?}", plan.mode),
                "epistemic_literal_count": plan.epistemic_literals.len(),
                "required_phases": plan.required_phases.iter().map(|phase| format!("{:?}", phase)).collect::<Vec<_>>(),
                "required_kernel_phases": plan.required_kernel_phases.iter().map(|phase| format!("{:?}", phase)).collect::<Vec<_>>(),
                "required_buffers": plan.required_buffers.iter().map(|buffer| format!("{:?}", buffer)).collect::<Vec<_>>(),
                "reductions": reductions,
                "tuple_membership_bindings": tuple_membership,
                "solver_contract": {
                    "assumption_count": plan.solver_contract.assumption_bindings.len(),
                    "required_capabilities": plan.solver_contract.required_capabilities.iter().map(|cap| format!("{:?}", cap)).collect::<Vec<_>>(),
                    "required_statuses": plan.solver_contract.required_statuses.iter().map(|status| format!("{:?}", status)).collect::<Vec<_>>(),
                },
                "execution_backend": epistemic_execution_backend_label(plan.execution_backend),
                "fallback_policy": epistemic_fallback_policy_label(plan.fallback_policy),
            })
        }
        Err(err) => serde_json::json!({
            "status": "error",
            "error": err.to_string(),
        }),
    };

    let executable_plan = match xlog_logic::epistemic::compile_epistemic_gpu_execution(program) {
        Ok(plan) => serde_json::json!({
            "status": "ok",
            "relation_id_count": plan.relation_ids.len(),
            "reduced_runtime_sccs": plan.reduced_runtime_plan.sccs.len(),
            "reduced_runtime_est_memory_peak": plan.reduced_runtime_plan.est_memory_peak,
            "gpu_plan_literal_count": plan.gpu_plan.epistemic_literals.len(),
            "execution_backend": epistemic_execution_backend_label(plan.gpu_plan.execution_backend),
            "fallback_policy": epistemic_fallback_policy_label(plan.gpu_plan.fallback_policy),
        }),
        Err(err) => serde_json::json!({
            "status": "error",
            "error": err.to_string(),
        }),
    };

    serde_json::json!({
        "eir": eir,
        "gpu_plan": gpu_plan,
        "executable_plan": executable_plan,
    })
}

fn program_has_epistemic_literals(program: &Program) -> bool {
    program.rules.iter().any(|rule| {
        rule.body
            .iter()
            .any(|lit| matches!(lit, BodyLiteral::Epistemic(_)))
    }) || program.constraints.iter().any(|constraint| {
        constraint
            .body
            .iter()
            .any(|lit| matches!(lit, BodyLiteral::Epistemic(_)))
    })
}

fn eir_atom_json(atom: &xlog_ir::EirAtom) -> serde_json::Value {
    serde_json::json!({
        "predicate": &atom.predicate,
        "arity": atom.arity,
        "terms": atom.terms.iter().map(eir_term_label).collect::<Vec<_>>(),
    })
}

fn eir_epistemic_literal_json(lit: &xlog_ir::EirEpistemicLiteral) -> serde_json::Value {
    serde_json::json!({
        "op": format!("{:?}", lit.op),
        "negated": lit.negated,
        "atom": eir_atom_json(&lit.atom),
    })
}

fn epistemic_execution_backend_label(backend: xlog_ir::EpistemicExecutionBackend) -> &'static str {
    match backend {
        xlog_ir::EpistemicExecutionBackend::Gpu => "gpu",
    }
}

fn epistemic_fallback_policy_label(policy: xlog_ir::EpistemicFallbackPolicy) -> &'static str {
    match policy {
        xlog_ir::EpistemicFallbackPolicy::RejectUnsupported => "reject_unsupported",
    }
}

fn eir_term_label(term: &EirTerm) -> String {
    match term {
        EirTerm::Variable(name) => name.clone(),
        EirTerm::Anonymous => "_".to_string(),
        EirTerm::Integer(value) => value.to_string(),
        EirTerm::FloatBits(bits) => f64::from_bits(*bits).to_string(),
        EirTerm::String(value) => value.clone(),
        EirTerm::Symbol(id) => symbol::resolve(*id),
        EirTerm::List(items) => format!(
            "[{}]",
            items
                .iter()
                .map(eir_term_label)
                .collect::<Vec<_>>()
                .join(", ")
        ),
        EirTerm::Cons { head, tail } => {
            format!("{}|{}", eir_term_label(head), eir_term_label(tail))
        }
        EirTerm::Compound { functor, args } => format!(
            "{}({})",
            functor,
            args.iter()
                .map(eir_term_label)
                .collect::<Vec<_>>()
                .join(", ")
        ),
        EirTerm::PredRef(name) => name.clone(),
        EirTerm::Aggregate { op, variable } => format!("{}({})", op, variable),
    }
}

fn explain_aggregate_lifting(program: &Program) -> Result<Vec<AggregateLiftReport>> {
    let has_probabilistic_source =
        !program.prob_facts.is_empty() || !program.annotated_disjunctions.is_empty();
    let has_aggregate_rule = program.proper_rules().any(|rule| rule.has_aggregation());
    if !(has_probabilistic_source && has_aggregate_rule) {
        return Ok(Vec::new());
    }
    Ok(xlog_prob::provenance::extract_from_program(program)?.aggregate_lifting)
}

fn print_explain_text(report: &ExplainReport) {
    println!("parse:");
    println!("  statements: {}", report.parse_stats.statement_count);
    println!("ast:");
    println!("  rules: {}", report.program.rules.len());
    println!("  queries: {}", report.program.queries.len());
    println!("stratification:");
    println!("  status: {}", report.stratification_status);
    if let Some(reason) = &report.stratification_reason {
        println!("  reason: {reason}");
    }
    println!("  strata: {}", report.stratification_count);
    println!("rir:");
    println!("  status: {}", report.rir_status);
    if let Some(reason) = &report.rir_reason {
        println!("  reason: {reason}");
    }
    println!("  sccs: {}", report.rir_sccs);
    println!("optimizer:");
    println!("  status: {}", report.optimizer_status);
    if let Some(reason) = &report.optimizer_reason {
        println!("  reason: {reason}");
    }
    println!("  est_memory_peak: {}", report.optimizer_memory_peak);
    println!("wcoj:");
    println!("  status: {}", report.wcoj_status);
    if let Some(reason) = &report.wcoj_reason {
        println!("  reason: {reason}");
    }
    println!("epistemic:");
    for section in ["eir", "gpu_plan", "executable_plan"] {
        let analysis = report.epistemic.get(section);
        let status = analysis
            .and_then(|value| value.get("status"))
            .and_then(serde_json::Value::as_str)
            .unwrap_or("not_available");
        println!("  {section}: {status}");
        if let Some(reason) = analysis
            .and_then(|value| value.get("reason"))
            .and_then(serde_json::Value::as_str)
        {
            println!("  {section}_reason: {reason}");
        }
    }
    print_magic_text(&report.magic_sets);
    if report.aggregate_lifting_status != "ok" || !report.aggregate_lifting.is_empty() {
        println!("aggregate_lifting:");
        println!("  status: {}", report.aggregate_lifting_status);
        if let Some(reason) = &report.aggregate_lifting_reason {
            println!("  reason: {reason}");
        }
        for entry in &report.aggregate_lifting {
            println!(
                "  - predicate: {} operator: {} status: {} domain: {} uncertain: {} cap: {}",
                entry.predicate,
                entry.operator,
                entry.status.as_str(),
                entry.domain_size,
                entry.uncertain_rows,
                entry.cap
            );
        }
    }
    if !report.rule_provenance.is_empty() {
        println!("rule_provenance:");
        for entry in &report.rule_provenance {
            println!(
                "  - id: {} source_kind: {} head: {}",
                entry.rule_id,
                entry.source_kind.as_str(),
                entry.head
            );
        }
    }
    if !report.proof_traces.is_empty() {
        println!("proof_traces:");
        for entry in &report.proof_traces {
            println!(
                "  - query: {} rules: {} source_facts: {}",
                entry.query,
                entry.rule_ids.len(),
                entry.source_facts.len()
            );
        }
    }
    println!("generated_rule_diagnostics:");
    println!("  status: {}", report.generated_rule_diagnostics_status);
    if let Some(reason) = &report.generated_rule_diagnostics_reason {
        println!("  reason: {reason}");
    }
    println!("  rules: {}", report.generated_rule_diagnostics.len());
}

fn print_magic_text(report: &MagicSetReport) {
    println!("magic_sets:");
    println!("  status: {}", magic_status_label(report.status));
    if !report.adorned_predicates.is_empty() {
        println!("  adorned_predicates:");
        for pred in &report.adorned_predicates {
            println!("    - {}", pred);
        }
    }
    if !report.generated_predicates.is_empty() {
        println!("  generated_predicates:");
        for pred in &report.generated_predicates {
            println!("    - {}", pred);
        }
    }
    if !report.declined_reasons.is_empty() {
        println!("  declined_reasons:");
        for reason in &report.declined_reasons {
            println!("    - {}", reason);
        }
    }
}

fn print_explain_json(report: &ExplainReport) {
    println!("{{");
    println!("  \"parse\": {{");
    println!(
        "    \"statements\": {},",
        report.parse_stats.statement_count
    );
    println!("    \"cache_hits\": {},", report.parse_stats.hits);
    println!("    \"cache_misses\": {}", report.parse_stats.misses);
    println!("  }},");
    println!("  \"ast\": {{");
    println!("    \"rules\": {},", report.program.rules.len());
    println!("    \"queries\": {},", report.program.queries.len());
    println!(
        "    \"prob_queries\": {}",
        report.program.prob_queries.len()
    );
    println!("  }},");
    println!("  \"stratification\": {{");
    println!(
        "    \"status\": \"{}\",",
        json_escape(&report.stratification_status)
    );
    println!(
        "    \"reason\": {},",
        json_optional_string(report.stratification_reason.as_deref())
    );
    println!("    \"strata\": {}", report.stratification_count);
    println!("  }},");
    println!("  \"rir\": {{");
    println!("    \"status\": \"{}\",", json_escape(&report.rir_status));
    println!(
        "    \"reason\": {},",
        json_optional_string(report.rir_reason.as_deref())
    );
    println!("    \"sccs\": {}", report.rir_sccs);
    println!("  }},");
    println!("  \"optimizer\": {{");
    println!(
        "    \"status\": \"{}\",",
        json_escape(&report.optimizer_status)
    );
    println!(
        "    \"reason\": {},",
        json_optional_string(report.optimizer_reason.as_deref())
    );
    println!("    \"est_memory_peak\": {}", report.optimizer_memory_peak);
    println!("  }},");
    println!("  \"wcoj\": {{");
    println!("    \"status\": \"{}\",", json_escape(&report.wcoj_status));
    println!(
        "    \"reason\": {}",
        json_optional_string(report.wcoj_reason.as_deref())
    );
    println!("  }},");
    println!("  \"epistemic\": {},", report.epistemic);
    println!("  \"magic_sets\": {{");
    println!(
        "    \"status\": \"{}\",",
        json_escape(magic_status_label(report.magic_sets.status))
    );
    println!(
        "    \"adorned_predicates\": {},",
        json_string_array(&report.magic_sets.adorned_predicates)
    );
    println!(
        "    \"generated_predicates\": {},",
        json_string_array(&report.magic_sets.generated_predicates)
    );
    println!(
        "    \"declined_reasons\": {}",
        json_string_array(&report.magic_sets.declined_reasons)
    );
    println!("  }},");
    println!("  \"probability\": {{");
    println!(
        "    \"engine\": \"{}\",",
        match report.program.prob_engine() {
            ProbEngine::ExactDdnnf => "exact_ddnnf",
            ProbEngine::Mc => "mc",
        }
    );
    println!(
        "    \"aggregate_lifting_status\": \"{}\",",
        json_escape(&report.aggregate_lifting_status)
    );
    println!(
        "    \"aggregate_lifting_reason\": {},",
        json_optional_string(report.aggregate_lifting_reason.as_deref())
    );
    println!(
        "    \"aggregate_lifting_count\": {}",
        report.aggregate_lifting.len()
    );
    println!("  }},");
    println!("  \"aggregate_lifting\": [");
    for (idx, entry) in report.aggregate_lifting.iter().enumerate() {
        let suffix = if idx + 1 == report.aggregate_lifting.len() {
            ""
        } else {
            ","
        };
        println!("    {{");
        println!(
            "      \"predicate\": \"{}\",",
            json_escape(&entry.predicate)
        );
        println!(
            "      \"group_key\": {},",
            json_value_array(&entry.group_key)
        );
        println!("      \"operator\": \"{}\",", json_escape(&entry.operator));
        println!(
            "      \"finite_domain_source\": \"{}\",",
            json_escape(&entry.finite_domain_source)
        );
        println!(
            "      \"deterministic_rows\": {},",
            entry.deterministic_rows
        );
        println!("      \"uncertain_rows\": {},", entry.uncertain_rows);
        println!("      \"domain_size\": {},", entry.domain_size);
        println!("      \"cap\": {},", entry.cap);
        println!("      \"status\": \"{}\",", entry.status.as_str());
        println!("      \"reason\": \"{}\",", json_escape(&entry.reason));
        println!("      \"naive_outcomes\": {},", entry.naive_outcomes);
        println!(
            "      \"dynamic_programming_states\": {}",
            entry.dynamic_programming_states
        );
        println!("    }}{}", suffix);
    }
    println!("  ],");
    print_rule_provenance_json(&report.rule_provenance);
    println!(",");
    print_proof_traces_json(&report.proof_traces);
    println!(",");
    println!(
        "  \"generated_rule_diagnostics_status\": \"{}\",",
        json_escape(&report.generated_rule_diagnostics_status)
    );
    println!(
        "  \"generated_rule_diagnostics_reason\": {},",
        json_optional_string(report.generated_rule_diagnostics_reason.as_deref())
    );
    print_generated_rule_diagnostics_json(&report.generated_rule_diagnostics);
    println!("}}");
}

fn print_rule_provenance_json(entries: &[RuleProvenance]) {
    println!("  \"rule_provenance\": [");
    for (idx, entry) in entries.iter().enumerate() {
        let suffix = if idx + 1 == entries.len() { "" } else { "," };
        println!("    {{");
        println!("      \"rule_id\": \"{}\",", json_escape(&entry.rule_id));
        println!("      \"head\": \"{}\",", json_escape(&entry.head));
        println!(
            "      \"source_kind\": \"{}\",",
            json_escape(entry.source_kind.as_str())
        );
        println!(
            "      \"source_span\": {},",
            json_optional_string(entry.source_span.as_deref())
        );
        println!(
            "      \"generation_trace_hash\": {},",
            json_optional_string(entry.generation_trace_hash.as_deref())
        );
        println!(
            "      \"support_relation_ids\": {},",
            json_string_array(&entry.support_relation_ids)
        );
        println!(
            "      \"counterexample_relation_ids\": {}",
            json_string_array(&entry.counterexample_relation_ids)
        );
        println!("    }}{}", suffix);
    }
    println!("  ]");
}

fn print_proof_traces_json(entries: &[QueryProofTrace]) {
    println!("  \"proof_traces\": [");
    for (idx, entry) in entries.iter().enumerate() {
        let suffix = if idx + 1 == entries.len() { "" } else { "," };
        println!("    {{");
        println!("      \"query_id\": \"{}\",", json_escape(&entry.query_id));
        println!("      \"query\": \"{}\",", json_escape(&entry.query));
        println!(
            "      \"answer_relation\": \"{}\",",
            json_escape(&entry.answer_relation)
        );
        println!(
            "      \"rule_ids\": {},",
            json_string_array(&entry.rule_ids)
        );
        println!(
            "      \"source_facts\": {},",
            json_string_array(&entry.source_facts)
        );
        println!(
            "      \"rejected_alternatives\": {}",
            json_string_array(&entry.rejected_alternatives)
        );
        println!("    }}{}", suffix);
    }
    println!("  ]");
}

fn print_magic_dot(report: &MagicSetReport) {
    println!("digraph xlog_magic_sets {{");
    println!(
        "  status [label=\"status: {}\"];",
        magic_status_label(report.status)
    );
    for pred in &report.generated_predicates {
        println!("  \"{}\" [shape=box];", dot_escape(pred));
    }
    for pred in &report.adorned_predicates {
        println!("  \"{}\" [shape=ellipse];", dot_escape(pred));
    }
    for (index, reason) in report.declined_reasons.iter().enumerate() {
        println!(
            "  reason_{index} [shape=note,label=\"reason: {}\"];",
            dot_escape(reason)
        );
        println!("  status -> reason_{index} [style=dashed];");
    }
    println!("}}");
}

fn magic_status_label(status: MagicSetStatus) -> &'static str {
    match status {
        MagicSetStatus::Disabled => "disabled",
        MagicSetStatus::Applied => "applied",
        MagicSetStatus::Declined => "declined",
    }
}

fn json_string_array(items: &[String]) -> String {
    let values = items
        .iter()
        .map(|item| format!("\"{}\"", json_escape(item)))
        .collect::<Vec<_>>()
        .join(", ");
    format!("[{}]", values)
}

fn json_value_array(items: &[Value]) -> String {
    let values = items.iter().map(json_value).collect::<Vec<_>>().join(", ");
    format!("[{}]", values)
}

fn json_value(value: &Value) -> String {
    match value {
        Value::I64(v) => v.to_string(),
        Value::F64(bits) => {
            let v = f64::from_bits(*bits);
            if v.is_finite() {
                v.to_string()
            } else {
                format!("\"{}\"", json_escape(&v.to_string()))
            }
        }
        Value::Symbol(id) => format!("\"{}\"", json_escape(&symbol::resolve(*id))),
        Value::String(s) => format!("\"{}\"", json_escape(s)),
    }
}

fn json_optional_string(value: Option<&str>) -> String {
    match value {
        Some(value) => format!("\"{}\"", json_escape(value)),
        None => "null".to_string(),
    }
}

fn json_escape(value: &str) -> String {
    let serialized = serde_json::to_string(value).expect("serializing a string cannot fail");
    serialized[1..serialized.len() - 1].to_string()
}

fn dot_escape(value: &str) -> String {
    value.replace('\\', "\\\\").replace('"', "\\\"")
}

fn memory_budget_bytes(memory_mb: u64) -> Result<u64> {
    memory_mb.checked_mul(1024 * 1024).ok_or_else(|| {
        XlogError::Execution(format!("memory budget {memory_mb} MiB overflows bytes"))
    })
}

fn make_provider(device: usize, memory_mb: u64) -> Result<Arc<CudaKernelProvider>> {
    use xlog_cuda::device_runtime::{
        AsyncCudaResource, DeviceMemoryResource, GlobalDeviceBudget, StreamPool, XlogDeviceRuntime,
    };
    let memory_bytes = memory_budget_bytes(memory_mb)?;
    let runtime_memory_bytes = usize::try_from(memory_bytes).map_err(|_| {
        XlogError::Execution(format!(
            "memory budget {memory_mb} MiB cannot be represented on this platform"
        ))
    })?;
    let device = Arc::new(CudaDevice::new(device)?);
    // Runtime-backed memory manager: the recorded GPU primitives (WCOJ
    // triangle/4-cycle/k-clique, Free Join, factorized delta) require a
    // DeviceBlock-backed allocation routed through an XlogDeviceRuntime.
    // The plain `GpuMemoryManager::new` path leaves `memory().runtime()`
    // == None, so those dispatches silently fall back to binary joins.
    let pool = Arc::new(StreamPool::with_defaults(Arc::clone(&device)));
    let async_resource: Box<dyn DeviceMemoryResource + Send + Sync> = Box::new(
        AsyncCudaResource::new(Arc::clone(&device), 0, Arc::clone(&pool)),
    );
    let budget_resource: Box<dyn DeviceMemoryResource + Send + Sync> = Box::new(
        GlobalDeviceBudget::new(async_resource, runtime_memory_bytes),
    );
    let runtime = Arc::new(XlogDeviceRuntime::with_resource(
        Arc::clone(&device),
        0,
        Arc::clone(&pool),
        budget_resource,
    ));
    let memory = Arc::new(GpuMemoryManager::with_runtime(
        Arc::clone(&device),
        MemoryBudget::with_limit(memory_bytes),
        Arc::clone(&runtime),
    ));
    Ok(Arc::new(CudaKernelProvider::with_runtime(device, memory)?))
}

fn parse_inputs(inputs: &[String]) -> Result<HashMap<String, PathBuf>> {
    let mut out = HashMap::new();
    for entry in inputs {
        let (name, path) = entry.split_once('=').ok_or_else(|| {
            XlogError::Execution(format!("Invalid --input '{}', expected rel=path", entry))
        })?;
        out.insert(name.to_string(), PathBuf::from(path));
    }
    Ok(out)
}

fn run_deterministic(args: RunArgs) -> Result<()> {
    if args.wcoj {
        // Force the WCOJ dispatch gates (default RuntimeConfig consults these
        // env vars; see xlog_core::RuntimeConfig::wcoj_triangle_dispatch).
        std::env::set_var("XLOG_USE_WCOJ_TRIANGLE_U32", "1");
        std::env::set_var("XLOG_USE_WCOJ_4CYCLE", "1");
        // This only sets the gate. Whether a WCOJ kernel actually dispatched
        // (vs. silently falling back to binary joins) is reported post-run
        // via --stats `wcoj.triangle_dispatch` / `wcoj.four_cycle_dispatch`.
        eprintln!("WCOJ dispatch gates set (triangle + 4-cycle); run with --stats to confirm a kernel fired");
    }
    let provider = make_provider(args.device, args.memory_mb)?;
    let source = std::fs::read_to_string(&args.source).map_err(|e| {
        XlogError::Execution(format!("Failed to read {}: {}", args.source.display(), e))
    })?;

    let parsed = xlog_logic::parse_program(&source)?;
    let resolved = resolve_program_imports(
        parsed,
        &args.source,
        args.module_path.clone(),
        ModuleMergeErrorKind::Compilation,
    )?;
    let program = LogicProgram::compile_program(resolved)?;
    let mut inputs = HashMap::new();
    for (name, path) in parse_inputs(&args.input)? {
        let buf = provider.read_arrow_ipc_stream_file(&path)?;
        inputs.insert(name, buf);
    }

    let result = program.evaluate_with_options(provider.clone(), inputs, args.stats)?;

    // Dump the compiled epistemic execution plan after a successful GPU run, so
    // the JSON corresponds to a real accepted hot-path execution.
    if let Some(plan_path) = &args.epistemic_plan_json {
        match program.epistemic_plan_json() {
            Some(json) => {
                std::fs::write(plan_path, json).map_err(|e| {
                    XlogError::Execution(format!(
                        "Failed to write epistemic plan JSON {}: {}",
                        plan_path.display(),
                        e
                    ))
                })?;
                eprintln!("epistemic plan dumped to {}", plan_path.display());
            }
            None => {
                eprintln!(
                    "note: --epistemic-plan-json given but program has no epistemic literals; no plan dumped"
                );
            }
        }
    }

    // Emit query results
    emit_logic_results(
        provider.as_ref(),
        &result.queries,
        args.output,
        args.output_dir.as_deref(),
    )?;

    // Emit stats if requested
    if args.stats {
        if let Some(stats) = result.stats {
            if args.wcoj {
                // Honest confirmation: did a WCOJ kernel actually fire, or did
                // the run silently fall back to binary joins? Count every WCOJ
                // dispatch kind — fused triangle COUNTING, for instance, routes
                // through the group-by-fusion kernel, not the triangle hook.
                let fired = stats.wcoj_triangle_dispatch_count
                    + stats.wcoj_4cycle_dispatch_count
                    + stats.wcoj_groupby_fusion_dispatch_count
                    + stats.free_join_dispatch_count
                    + stats.factorized_delta_dispatch_count;
                if fired > 0 {
                    eprintln!(
                        "WCOJ kernels dispatched: triangle {}, 4-cycle {}, groupby-fusion {}, free-join {}, factorized-delta {} (declines {})",
                        stats.wcoj_triangle_dispatch_count,
                        stats.wcoj_4cycle_dispatch_count,
                        stats.wcoj_groupby_fusion_dispatch_count,
                        stats.free_join_dispatch_count,
                        stats.factorized_delta_dispatch_count,
                        stats.wcoj_error_decline_count,
                    );
                } else {
                    eprintln!(
                        "WARNING: --wcoj set but no WCOJ kernel dispatched (declines {}); the run fell back to binary joins",
                        stats.wcoj_error_decline_count,
                    );
                }
            }
            let stats_output = match args.stats_format {
                StatsFormat::Human => stats.format_human(),
                StatsFormat::Json => stats.format_json(),
            };
            eprintln!("{}", stats_output);
        }
        // Symbol table statistics
        eprintln!(
            "Symbols: {} interned ({} bytes)",
            symbol::count(),
            symbol::memory_usage()
        );
    }

    Ok(())
}

fn run_probabilistic(args: ProbArgs) -> Result<()> {
    #[cfg(not(feature = "host-io"))]
    {
        let _ = args;
        return Err(XlogError::Execution(
            "Host output is disabled (feature \"host-io\" is OFF). Use device-resident APIs (DLPack) or rebuild with --features host-io.".to_string(),
        ));
    }

    #[cfg(feature = "host-io")]
    {
        let source = std::fs::read_to_string(&args.source).map_err(|e| {
            XlogError::Execution(format!("Failed to read {}: {}", args.source.display(), e))
        })?;
        let parsed_program = parse_program(&source)?;
        let program = resolve_program_imports(
            parsed_program,
            &args.source,
            args.module_path.clone(),
            ModuleMergeErrorKind::Execution,
        )?;

        let mut config = GpuConfig::default();
        config.device_ordinal = args.device;
        config.memory_bytes = memory_budget_bytes(args.memory_mb)?;

        match resolve_prob_engine(&args, &program) {
            ProbEngineCli::ExactDdnnf => {
                let prog = ExactDdnnfProgram::compile_from_program(&program, config)?;
                let result = prog.evaluate()?;
                emit_prob_exact(result, args.output, args.output_dir.as_deref())
            }
            ProbEngineCli::Mc => {
                let prog = McProgram::compile_from_program(&program, config)?;
                let mut cfg = McEvalConfig::from_directives(&program.directives)?;
                apply_mc_cli_overrides(&args, &mut cfg)?;
                // `evaluate` runs the GPU-native device hot loop and then
                // materializes the result on the host (downloads the final
                // query/evidence counts after the loop) so the CLI can print
                // probabilities and confidence intervals. The hot loop itself is
                // zero-host; this final download is host-result materialization,
                // not a hot-loop transfer. Device-resident consumers that want to
                // keep counts on the GPU use `McProgram::evaluate_gpu_device`.
                let result = prog.evaluate(cfg)?;
                emit_prob_mc(result, args.output, args.output_dir.as_deref())
            }
        }
    }
}

#[cfg(feature = "host-io")]
fn resolve_prob_engine(args: &ProbArgs, program: &Program) -> ProbEngineCli {
    args.prob_engine
        .unwrap_or_else(|| match program.directives.prob_engine_or_default() {
            ProbEngine::ExactDdnnf => ProbEngineCli::ExactDdnnf,
            ProbEngine::Mc => ProbEngineCli::Mc,
        })
}

#[cfg(feature = "host-io")]
fn apply_mc_cli_overrides(args: &ProbArgs, cfg: &mut McEvalConfig) -> Result<()> {
    if let Some(samples) = args.samples {
        cfg.samples = samples;
    }
    if let Some(seed) = args.seed {
        cfg.seed = seed;
    }
    if let Some(confidence) = args.confidence {
        cfg.confidence = confidence;
    }
    if let Some(iterations) = args.prob_max_nonmonotone_iterations {
        cfg.max_nonmonotone_iterations = iterations;
    }
    if let Some(method) = args.prob_method {
        cfg.sampling_method = Some(match method {
            ProbMethodCli::Rejection => McSamplingMethod::Rejection,
            ProbMethodCli::EvidenceClamping => McSamplingMethod::EvidenceClamping,
        });
    }
    cfg.allow_cpu_oracle_fallback = args.allow_cpu_oracle;
    cfg.validate()
}

fn emit_logic_results(
    provider: &CudaKernelProvider,
    queries: &[xlog_gpu::logic::LogicQueryResult],
    format: OutputFormat,
    output_dir: Option<&Path>,
) -> Result<()> {
    for (i, q) in queries.iter().enumerate() {
        if q.buffer.schema().arity() == 0 && matches!(format, OutputFormat::Pretty) {
            println!(
                "{}\nrows: {}",
                q.relation_name,
                provider.device_row_count(&q.buffer)?
            );
            continue;
        }
        let batch = provider.to_arrow_record_batch(&q.buffer)?;
        match format {
            OutputFormat::Pretty => {
                let formatted = pretty_format_batches(&[batch])
                    .map_err(|e| XlogError::Execution(format!("Pretty print failed: {}", e)))?;
                println!("{}\n{}", q.relation_name, formatted);
            }
            OutputFormat::Csv => {
                let mut out = Vec::new();
                {
                    let mut writer = WriterBuilder::new().build(&mut out);
                    writer
                        .write(&batch)
                        .map_err(|e| XlogError::Execution(format!("CSV write failed: {}", e)))?;
                }
                println!("{}\n{}", q.relation_name, String::from_utf8_lossy(&out));
            }
            OutputFormat::Arrow => {
                let dir = output_dir.unwrap_or_else(|| Path::new("."));
                let path = dir.join(format!("query_{}.arrow", i));
                provider.write_arrow_ipc_stream_file(&q.buffer, &path)?;
                println!("wrote {}", path.display());
            }
        }
    }
    Ok(())
}

#[cfg(feature = "host-io")]
fn emit_prob_exact(
    result: xlog_prob::exact::ExactResult,
    format: ProbOutputFormat,
    output_dir: Option<&Path>,
) -> Result<()> {
    if matches!(format, ProbOutputFormat::Json) {
        print_prob_exact_json(result);
        return Ok(());
    }

    let mut atoms = Vec::new();
    let mut probs = Vec::new();
    let mut log_probs = Vec::new();
    for q in result.query_probs {
        atoms.push(atom_to_string(&q.atom));
        probs.push(q.prob);
        log_probs.push(q.log_prob);
    }

    let batch = arrow::record_batch::RecordBatch::try_from_iter(vec![
        (
            "atom",
            Arc::new(arrow::array::StringArray::from(atoms)) as Arc<dyn arrow::array::Array>,
        ),
        (
            "prob",
            Arc::new(arrow::array::Float64Array::from(probs)) as Arc<dyn arrow::array::Array>,
        ),
        (
            "log_prob",
            Arc::new(arrow::array::Float64Array::from(log_probs)) as Arc<dyn arrow::array::Array>,
        ),
    ])
    .map_err(|e| XlogError::Execution(format!("Failed to build prob batch: {}", e)))?;

    emit_batch(
        "prob",
        &batch,
        prob_output_as_batch_format(format),
        output_dir,
    )
}

#[cfg(feature = "host-io")]
fn emit_prob_mc(
    result: xlog_prob::mc::McResult,
    format: ProbOutputFormat,
    output_dir: Option<&Path>,
) -> Result<()> {
    if matches!(format, ProbOutputFormat::Json) {
        print_prob_mc_json(result);
        return Ok(());
    }

    let total_samples = result.total_samples as u64;
    let evidence_samples = result.evidence_samples as u64;
    let seed = result.seed;
    let confidence = result.confidence;
    let sampling_method = result.sampling_method.as_str().to_string();
    let mc_engine = result.engine.as_str().to_string();

    let mut atoms = Vec::new();
    let mut probs = Vec::new();
    let mut log_probs = Vec::new();
    let mut stderr = Vec::new();
    let mut ci_low = Vec::new();
    let mut ci_high = Vec::new();
    let mut total_samples_col = Vec::new();
    let mut evidence_samples_col = Vec::new();
    let mut seed_col = Vec::new();
    let mut confidence_col = Vec::new();
    let mut sampling_method_col = Vec::new();
    let mut mc_engine_col = Vec::new();
    for q in result.query_estimates {
        atoms.push(atom_to_string(&q.atom));
        probs.push(q.prob);
        log_probs.push(q.log_prob);
        stderr.push(q.stderr);
        ci_low.push(q.ci_low);
        ci_high.push(q.ci_high);
        total_samples_col.push(total_samples);
        evidence_samples_col.push(evidence_samples);
        seed_col.push(seed);
        confidence_col.push(confidence);
        sampling_method_col.push(sampling_method.clone());
        mc_engine_col.push(mc_engine.clone());
    }

    let batch = arrow::record_batch::RecordBatch::try_from_iter(vec![
        (
            "atom",
            Arc::new(arrow::array::StringArray::from(atoms)) as Arc<dyn arrow::array::Array>,
        ),
        (
            "prob",
            Arc::new(arrow::array::Float64Array::from(probs)) as Arc<dyn arrow::array::Array>,
        ),
        (
            "log_prob",
            Arc::new(arrow::array::Float64Array::from(log_probs)) as Arc<dyn arrow::array::Array>,
        ),
        (
            "stderr",
            Arc::new(arrow::array::Float64Array::from(stderr)) as Arc<dyn arrow::array::Array>,
        ),
        (
            "ci_low",
            Arc::new(arrow::array::Float64Array::from(ci_low)) as Arc<dyn arrow::array::Array>,
        ),
        (
            "ci_high",
            Arc::new(arrow::array::Float64Array::from(ci_high)) as Arc<dyn arrow::array::Array>,
        ),
        (
            "total_samples",
            Arc::new(arrow::array::UInt64Array::from(total_samples_col))
                as Arc<dyn arrow::array::Array>,
        ),
        (
            "evidence_samples",
            Arc::new(arrow::array::UInt64Array::from(evidence_samples_col))
                as Arc<dyn arrow::array::Array>,
        ),
        (
            "seed",
            Arc::new(arrow::array::UInt64Array::from(seed_col)) as Arc<dyn arrow::array::Array>,
        ),
        (
            "confidence",
            Arc::new(arrow::array::Float64Array::from(confidence_col))
                as Arc<dyn arrow::array::Array>,
        ),
        (
            "sampling_method",
            Arc::new(arrow::array::StringArray::from(sampling_method_col))
                as Arc<dyn arrow::array::Array>,
        ),
        (
            "mc_engine",
            Arc::new(arrow::array::StringArray::from(mc_engine_col))
                as Arc<dyn arrow::array::Array>,
        ),
    ])
    .map_err(|e| XlogError::Execution(format!("Failed to build mc batch: {}", e)))?;

    emit_batch(
        "prob",
        &batch,
        prob_output_as_batch_format(format),
        output_dir,
    )
}

#[cfg(feature = "host-io")]
fn prob_output_as_batch_format(format: ProbOutputFormat) -> OutputFormat {
    match format {
        ProbOutputFormat::Pretty => OutputFormat::Pretty,
        ProbOutputFormat::Csv => OutputFormat::Csv,
        ProbOutputFormat::Arrow => OutputFormat::Arrow,
        ProbOutputFormat::Json => unreachable!("json output is handled before batch emission"),
    }
}

#[cfg(feature = "host-io")]
fn print_prob_exact_json(result: xlog_prob::exact::ExactResult) {
    println!("{{");
    println!("  \"engine\": \"exact_ddnnf\",");
    println!("  \"queries\": [");
    let len = result.query_probs.len();
    for (idx, q) in result.query_probs.into_iter().enumerate() {
        let suffix = if idx + 1 == len { "" } else { "," };
        println!("    {{");
        println!(
            "      \"atom\": \"{}\",",
            json_escape(&atom_to_string(&q.atom))
        );
        println!("      \"prob\": {},", q.prob);
        println!("      \"log_prob\": {}", q.log_prob);
        println!("    }}{}", suffix);
    }
    println!("  ]");
    println!("}}");
}

#[cfg(feature = "host-io")]
fn print_prob_mc_json(result: xlog_prob::mc::McResult) {
    let total_samples = result.total_samples;
    let evidence_samples = result.evidence_samples;
    let seed = result.seed;
    let confidence = result.confidence;
    let sampling_method = result.sampling_method.as_str();
    let mc_engine = result.engine.as_str();
    println!("{{");
    println!("  \"engine\": \"mc\",");
    println!("  \"mc_engine\": \"{}\",", mc_engine);
    println!("  \"total_samples\": {},", total_samples);
    println!("  \"evidence_samples\": {},", evidence_samples);
    println!("  \"seed\": {},", seed);
    println!("  \"confidence\": {},", confidence);
    println!("  \"sampling_method\": \"{}\",", sampling_method);
    println!("  \"queries\": [");
    let len = result.query_estimates.len();
    for (idx, q) in result.query_estimates.into_iter().enumerate() {
        let suffix = if idx + 1 == len { "" } else { "," };
        println!("    {{");
        println!(
            "      \"atom\": \"{}\",",
            json_escape(&atom_to_string(&q.atom))
        );
        println!("      \"prob\": {},", q.prob);
        println!("      \"log_prob\": {},", q.log_prob);
        println!("      \"stderr\": {},", q.stderr);
        println!("      \"ci_low\": {},", q.ci_low);
        println!("      \"ci_high\": {},", q.ci_high);
        println!("      \"total_samples\": {},", total_samples);
        println!("      \"evidence_samples\": {}", evidence_samples);
        println!("    }}{}", suffix);
    }
    println!("  ]");
    println!("}}");
}

#[cfg(feature = "host-io")]
fn emit_batch(
    name: &str,
    batch: &arrow::record_batch::RecordBatch,
    format: OutputFormat,
    output_dir: Option<&Path>,
) -> Result<()> {
    match format {
        OutputFormat::Pretty => {
            let formatted = pretty_format_batches(std::slice::from_ref(batch))
                .map_err(|e| XlogError::Execution(format!("Pretty print failed: {}", e)))?;
            println!("{}\n{}", name, formatted);
        }
        OutputFormat::Csv => {
            let mut out = Vec::new();
            {
                let mut writer = WriterBuilder::new().build(&mut out);
                writer
                    .write(batch)
                    .map_err(|e| XlogError::Execution(format!("CSV write failed: {}", e)))?;
            }
            println!("{}\n{}", name, String::from_utf8_lossy(&out));
        }
        OutputFormat::Arrow => {
            let dir = output_dir.unwrap_or_else(|| Path::new("."));
            let path = dir.join(format!("{}_prob.arrow", name));
            let mut out = Vec::new();
            let mut writer =
                arrow::ipc::writer::StreamWriter::try_new(&mut out, &batch.schema())
                    .map_err(|e| XlogError::Execution(format!("Arrow writer failed: {}", e)))?;
            writer
                .write(batch)
                .map_err(|e| XlogError::Execution(format!("Arrow write failed: {}", e)))?;
            writer
                .finish()
                .map_err(|e| XlogError::Execution(format!("Arrow finish failed: {}", e)))?;
            std::fs::write(&path, out)
                .map_err(|e| XlogError::Execution(format!("Arrow write file failed: {}", e)))?;
            println!("wrote {}", path.display());
        }
    }
    Ok(())
}

#[cfg(feature = "host-io")]
fn atom_to_string(atom: &xlog_prob::provenance::GroundAtom) -> String {
    use xlog_prob::provenance::Value;

    if atom.args.is_empty() {
        return format!("{}()", atom.predicate);
    }

    let mut out = String::new();
    out.push_str(&atom.predicate);
    out.push('(');
    for (i, arg) in atom.args.iter().enumerate() {
        if i != 0 {
            out.push_str(", ");
        }
        match arg {
            Value::I64(v) => out.push_str(&v.to_string()),
            Value::F64(bits) => out.push_str(&f64::from_bits(*bits).to_string()),
            Value::Symbol(sym) => out.push_str(&symbol::resolve(*sym)),
            Value::String(v) => out.push_str(v),
        }
    }
    out.push(')');
    out
}