dupehell 0.5.5

// DupeHell -- MIT License
//
// Synthetic multi-domain dataset generator for record linkage benchmarking.
// No liability for misuse.

use std::path::PathBuf;

use clap::Parser;

use dupehell_core::context::Context;
use dupehell_core::difficulty::estimate_difficulty;
use dupehell_core::pipeline::run_pipeline_with_progress;
use dupehell_core::schema::{
    build_pipeline_config, default_singleton_master_fraction, load_schema,
};

#[derive(Parser)]
#[command(
    name = "dupehell",
    version = env!("CARGO_PKG_VERSION"),
    about = "Synthetic record linkage dataset generator",
    long_about = "Generates realistic synthetic datasets with controlled duplicate rates, \
                  hard negatives, and noise profiles for benchmarking record linkage systems. \
                  Supports 40 domains (kyc, healthcare, ecommerce, gaming, ...), \
                  three difficulty levels, and outputs Arrow IPC or Parquet format."
)]
struct Cli {
    #[arg(
        long,
        default_value = "kyc",
        help = "Domain name (e.g. kyc, healthcare, gaming, ecommerce)"
    )]
    domain: String,

    #[arg(
        long,
        default_value_t = 1_000_000,
        help = "Number of base records to generate (minimum 10)"
    )]
    size: usize,

    #[arg(
        long,
        default_value_t = 42,
        help = "Random seed for deterministic reproducibility"
    )]
    seed: u64,

    #[arg(long, default_value = "medium", value_parser = clap::builder::PossibleValuesParser::new(["light", "medium", "hell"]), help = "Difficulty level: light, medium, or hell")]
    difficulty: String,

    #[arg(
        long,
        help = "Estimate difficulty and F1 score without generating data"
    )]
    estimate: bool,

    #[arg(
        long,
        help = "Pin the process to performance (P) cores only, on hybrid P-core/E-core CPUs (Windows only; no-op elsewhere or on non-hybrid CPUs)"
    )]
    pcore_only: bool,

    #[arg(
        long,
        help = "Output file format: parquet (default, ZSTD compressed) or ipc (Arrow IPC)"
    )]
    output_format: Option<String>,

    #[arg(long, action = clap::ArgAction::Count, help = "Shortcut for --output-format parquet (the default; kept for backward compatibility)")]
    parquet: u8,

    #[arg(
        long,
        default_value = ".",
        help = "Output directory (created automatically if missing)"
    )]
    output_dir: PathBuf,

    #[arg(
        long,
        default_value_t = 0.3,
        help = "Hard-negative scaling knob, not a literal fraction: actual count ~= n_duplicates * hard_neg_ratio * 0.125, where n_duplicates depends on --difficulty (light ~0.28*size, medium ~0.40*size, hell ~0.57*size); at medium + default 0.3 this is ~1.5% of size. Use --estimate to see the exact count first."
    )]
    hard_neg_ratio: f64,

    #[arg(
        long,
        help = "Fraction of masters with only one record (0.0 to 1.0). Defaults \
                to the chosen --difficulty tier's own value (0.50/0.30/0.10 \
                for light/medium/hell, see schema::default_singleton_master_fraction) \
                — pass this explicitly only to override that tier default."
    )]
    singleton_master_fraction: Option<f64>,

    #[arg(long, default_value = "en", value_parser = clap::builder::PossibleValuesParser::new(["en", "fr", "de", "es", "it", "pt"]), help = "Locale for pool data (en, fr, de, es, it, pt)")]
    locale: String,

    #[arg(
        long,
        default_value = "assets/pools",
        help = "Path to asset pools directory"
    )]
    pools_dir: PathBuf,

    #[arg(
        long,
        default_value = "schemas",
        help = "Path to schema JSON directory"
    )]
    schemas_dir: PathBuf,

    #[arg(
        long,
        help = "Generate property-graph output (nodes + edges) in addition to tabular data"
    )]
    graph: bool,

    #[arg(
        long,
        default_value = "parquet",
        value_parser = clap::builder::PossibleValuesParser::new(["ipc", "parquet"]),
        help = "Graph output format: parquet (default, ZSTD compressed) or ipc (requires --graph)"
    )]
    graph_format: String,

    #[arg(
        long,
        help = "Generate only this entity (e.g. --domain aviation --only-entity passenger). \
                --size then applies entirely to this entity instead of being split across the \
                domain's entities by weight. Other entities it references via FK are generated \
                only far enough to seed a plausible-looking identifier pool (capped, never \
                written to output); entities it doesn't reference are skipped entirely. Output \
                schema is narrowed to just this entity's own columns."
    )]
    only_entity: Option<String>,

    #[arg(
        long,
        help = "Generate internally as ceil(size / chunk-size) sequential chunks (each an \
                independently-seeded run, RAM-bounded like a --size chunk-size run) instead \
                of one single run, then assemble the results into the same single dataset/GT/ \
                graph files a non-chunked run would produce. record_id/master_id stay globally \
                contiguous across chunks. Use when --size is large enough that a single run \
                would exceed available RAM. No effect if omitted or >= --size."
    )]
    chunk_size: Option<usize>,
}

/// Conservative floor on bytes/record used only to flag genuinely tight
/// runs, not to project a peak RSS — actual usage depends heavily on
/// domain/difficulty/--graph and measured peaks stayed well under this
/// on every scale tried so far (see hunt2808.md). Not meant to be tightened
/// into an accurate estimator; it exists purely as a coarse tripwire.
const BYTES_PER_RECORD_FLOOR: usize = 150;

/// Best-effort warning if the requested `--size` looks likely to be tight on
/// available system RAM. Deliberately gives no GB projection — past
/// measurements showed the naive per-record estimate overshoots real usage
/// by a wide and inconsistent margin, so a numeric prediction here would
/// just be misleading. Advisory only — if system memory can't be read
/// (sandboxed/unusual environment), this silently does nothing rather than
/// block a run on a guess.
fn warn_if_memory_tight(size: usize) {
    let mut sys = sysinfo::System::new();
    sys.refresh_memory();
    let available = sys.available_memory();
    if available == 0 {
        return;
    }
    let floor = size as u64 * BYTES_PER_RECORD_FLOOR as u64;
    if floor > available {
        eprintln!(
            "Warning: --size {size} may be tight on available RAM (~{:.1} GB free). \
             If the run is unusually slow, it's likely swapping — consider a smaller \
             --size or splitting into multiple runs.",
            available as f64 / 1e9,
        );
    }
}

fn main() {
    env_logger::init();
    let cli = Cli::parse();

    if cli.pcore_only {
        if dupehell_core::cpu_affinity::pin_to_p_cores() {
            eprintln!("pcore-only: process pinned to P-cores");
        } else {
            eprintln!("pcore-only: no hybrid P-core/E-core topology detected, ignoring");
        }
    }

    let schema = match load_schema(&cli.domain, &cli.schemas_dir) {
        Ok(s) => s,
        Err(e) => {
            eprintln!("Error: {e}");
            std::process::exit(1);
        }
    };

    if cli.estimate {
        match estimate_difficulty(
            &cli.domain,
            cli.size,
            cli.seed,
            &cli.difficulty,
            cli.hard_neg_ratio,
            &schema,
        ) {
            Ok(report) => {
                println!("{}", serde_json::to_string_pretty(&report).unwrap());
            }
            Err(e) => {
                eprintln!("Estimation failed: {e}");
                std::process::exit(1);
            }
        }
        return;
    }

    if cli.size < 10 {
        eprintln!("Error: size must be >= 10, got {}", cli.size);
        std::process::exit(1);
    }
    const MAX_SIZE: usize = 500_000_000;
    if cli.size > MAX_SIZE {
        eprintln!(
            "Error: size must be <= {MAX_SIZE} (500M), got {}. \
             Larger runs risk exhausting memory in a single process; \
             split into multiple runs instead.",
            cli.size
        );
        std::process::exit(1);
    }
    if let Some(cs) = cli.chunk_size {
        if cs == 0 {
            eprintln!("Error: --chunk-size must be >= 1, got {cs}");
            std::process::exit(1);
        }
        // RAM is bounded per chunk (each is an independent pipeline run),
        // not by the total --size — that's the whole point of chunking.
        warn_if_memory_tight(cs.min(cli.size));
    } else {
        warn_if_memory_tight(cli.size);
    }

    let mut ctx = match Context::new(&cli.domain, &cli.locale, &cli.pools_dir.to_string_lossy()) {
        Ok(c) => c,
        Err(e) => {
            eprintln!("Error loading pools: {e}");
            std::process::exit(1);
        }
    };

    let effective_format = match &cli.output_format {
        Some(fmt) => {
            if cli.parquet > 0 && fmt != "parquet" {
                eprintln!(
                    "Warning: --parquet is ignored because --output-format {fmt} was also given"
                );
            }
            fmt.clone()
        }
        None => "parquet".to_string(),
    };

    if effective_format != "ipc" && effective_format != "parquet" {
        eprintln!("Error: output format must be 'ipc' or 'parquet', got '{effective_format}'");
        std::process::exit(1);
    }

    // Defaults to the chosen difficulty tier's own singleton fraction unless
    // explicitly overridden — previously this silently defaulted to 0.10
    // (the "hell" tier value) regardless of --difficulty, so every tier's
    // duplicate volume was effectively pinned at "hell" levels unless the
    // caller happened to also pass this flag by hand.
    let tier_default_singleton = default_singleton_master_fraction(&cli.difficulty);
    let singleton_master_fraction = match cli.singleton_master_fraction {
        Some(v) => {
            if (v - tier_default_singleton).abs() > f64::EPSILON {
                eprintln!(
                    "Warning: --singleton-master-fraction {v} overrides the '{}' tier's \
                     default ({tier_default_singleton}) — duplicate volume will differ from \
                     what --estimate reports unless it's given the same override.",
                    cli.difficulty
                );
            }
            v
        }
        None => tier_default_singleton,
    };

    let run_id = dupehell_core::schema::deterministic_run_id(
        &cli.domain,
        cli.size,
        cli.seed,
        &cli.difficulty,
        cli.hard_neg_ratio,
        singleton_master_fraction,
        &cli.locale,
        cli.only_entity.as_deref(),
        cli.chunk_size,
    );
    let config = match build_pipeline_config(
        &cli.domain,
        cli.size,
        cli.seed,
        &cli.difficulty,
        cli.hard_neg_ratio,
        singleton_master_fraction,
        &schema,
        &run_id,
        &effective_format,
        cli.graph,
        &cli.graph_format,
        cli.only_entity.as_deref(),
    ) {
        Ok(c) => c,
        Err(e) => {
            eprintln!("Error building config: {e}");
            std::process::exit(1);
        }
    };

    // `run_id` is deterministic (BUGS.md C14/C15 fixed it to hash every
    // parameter that affects the data), so a matching file only exists here
    // if this exact run was already generated before — warn instead of
    // silently overwriting it (BUGS.md C16). Advisory only: there's no
    // `--force` gate, since deliberately regenerating an identical run is a
    // normal, common thing to want to do.
    let dataset_ext = if effective_format == "parquet" {
        "parquet"
    } else {
        "ipc"
    };
    let dataset_path = cli.output_dir.join(format!("{run_id}.{dataset_ext}"));
    if dataset_path.exists() {
        eprintln!(
            "Warning: {} already exists and will be overwritten by this run.",
            dataset_path.display()
        );
    }

    ctx.enable_watermark(&config.domain, config.size, config.seed);

    eprintln!(
        "DupeHell v{} — {} domain, {} records [{}]",
        env!("CARGO_PKG_VERSION"),
        cli.domain.to_uppercase(),
        cli.size,
        cli.difficulty
    );
    eprintln!(
        "Entities: {} types, {} HN types",
        config.entity_plans.len(),
        config.hard_neg_types.len()
    );

    let t0 = std::time::Instant::now();
    // Progress line only for runs big enough that generation actually takes
    // a noticeable amount of wall time — printing/flushing on every batch of
    // a 10K-record run would just add noise, not information. Throttled to
    // ~1 update/second (not one per 500K-row batch) so it's readable instead
    // of scrolling past on runs with many small entities.
    const PROGRESS_MIN_SIZE: usize = 1_000_000;
    let target_size = cli.size;
    let mut last_print = std::time::Instant::now();
    let mut progress_cb = move |done: usize, total: usize| {
        if target_size < PROGRESS_MIN_SIZE {
            return;
        }
        let now = std::time::Instant::now();
        if now.duration_since(last_print).as_secs_f64() < 1.0 && done < total {
            return;
        }
        last_print = now;
        let pct = (done as f64 / total.max(1) as f64 * 100.0).min(100.0);
        eprint!("\r  Generating... {done}/{total} ({pct:.0}%)   ");
        let _ = std::io::Write::flush(&mut std::io::stderr());
    };
    let output_dir_str = cli.output_dir.to_string_lossy();
    let output = match cli.chunk_size {
        Some(cs) if cs > 0 && cs < cli.size => dupehell_core::pipeline::run_chunked(
            &ctx,
            &cli.domain,
            cli.size,
            cs,
            cli.seed,
            &cli.difficulty,
            cli.hard_neg_ratio,
            singleton_master_fraction,
            &schema,
            &run_id,
            &effective_format,
            cli.graph,
            &cli.graph_format,
            cli.only_entity.as_deref(),
            &output_dir_str,
            Some(&mut progress_cb),
        ),
        _ => run_pipeline_with_progress(&ctx, &config, &output_dir_str, Some(&mut progress_cb)),
    };
    let output = match output {
        Ok(o) => o,
        Err(e) => {
            eprintln!("Pipeline failed: {e}");
            std::process::exit(1);
        }
    };
    if cli.size >= PROGRESS_MIN_SIZE {
        eprintln!();
    }
    let elapsed = t0.elapsed().as_secs_f64();

    let n = output.stats.total_records;
    eprintln!(
        "\nDone in {:.3}s — {} records ({:.0} rec/s)",
        elapsed,
        n,
        n as f64 / elapsed
    );
    eprintln!(
        "  exact_dups={} fuzzy_dups={} hard_negs={} uniques={} masters={}",
        output.stats.exact_dups,
        output.stats.fuzzy_dups,
        output.stats.hard_negs,
        output.stats.uniques,
        output.stats.masters
    );
    eprintln!("  Dataset: {}", output.output_files[0]);
    eprintln!("  GT:      {}", output.gt_file);
    if let Some(nodes) = &output.nodes {
        eprintln!("  Nodes:  {nodes}");
    }
    if let Some(edges) = &output.edges {
        eprintln!("  Edges:  {edges}");
    }

    let id_cols: Vec<&str> = config
        .entity_plans
        .iter()
        .filter_map(|p| p.identifier_col.as_deref())
        .collect();
    if !id_cols.is_empty() {
        eprintln!(
            "  Note: {} are structural join keys (stable across all duplicates by design) — exclude them from ER match attributes, use record_id instead.",
            id_cols.join(", ")
        );
    }
}