ruwex 0.1.0

//! Drop-in replacement for wikiextractor's `wikiextractor` command.

use std::ffi::OsString;
use std::path::Path;
use std::sync::Arc;
use std::time::Instant;

use anyhow::{Context, ensure};
use clap::Parser;
use log::{LevelFilter, info};

use ruwex::expand::{LazyTemplateSource, TemplateSource};
use ruwex::{
    DocSink, ExtractorConfig, OutputFormat, PageSource, ShardedWriter, StdoutSink, TemplateDb,
    TitleIndex, default_workers, parse_size, render_page,
};

#[derive(Parser, Debug)]
#[command(
    name = "wikiextractor",
    about = "Extracts and cleans text from a Wikipedia database dump. \
             Rust rewrite (ruwex) of wikiextractor, CLI-compatible.",
    disable_version_flag = true
)]
struct Args {
    /// XML wiki dump file (.xml or .xml.bz2; multistream dumps with their
    /// index file alongside are processed in parallel), or - for stdin
    // Optional only so that a bare `-v` works like argparse's version action;
    // run() rejects a missing input otherwise.
    input: Option<String>,

    /// directory for extracted files (or '-' for dumping to stdout)
    #[arg(short, long, default_value = "text")]
    output: String,

    /// maximum bytes per output file (0 means no limit)
    #[arg(short, long, default_value = "1M", value_name = "n[KMG]")]
    bytes: String,

    /// compress output files using bzip
    #[arg(short, long)]
    compress: bool,

    /// write output in json format instead of the default <doc> format
    #[arg(long)]
    json: bool,

    /// produce HTML output, subsumes --links
    #[arg(long)]
    html: bool,

    /// preserve links
    #[arg(short, long)]
    links: bool,

    /// accepted namespaces, e.g. "Talk,Category" (-ns also works)
    #[arg(long, value_name = "ns1,ns2")]
    namespaces: Option<String>,

    /// use or create file containing templates
    #[arg(long, value_name = "TEMPLATES")]
    templates: Option<String>,

    /// do not expand templates
    #[arg(long)]
    no_templates: bool,

    /// produce HTML-safe output within <doc>...</doc>
    #[arg(long, value_name = "HTML_SAFE", num_args = 0..=1, default_missing_value = "true")]
    html_safe: Option<String>,

    /// number of extraction processes (default: number of cpus - 1)
    #[arg(long, value_name = "PROCESSES")]
    processes: Option<usize>,

    /// suppress reporting progress info
    #[arg(short, long)]
    quiet: bool,

    /// print debug info
    #[arg(long)]
    debug: bool,

    /// analyze a file containing a single article (debug option)
    #[arg(short, long)]
    article: bool,

    /// extract just the page with this exact title to stdout, using a title
    /// index built (once) beside the multistream dump
    #[arg(long, value_name = "TITLE")]
    title: Option<String>,

    /// print program version
    #[arg(short = 'v', long = "version")]
    version: bool,
}

/// argparse accepts `-ns`; clap short flags are single characters, so map it
/// to the long form before parsing.
fn preprocessed_args() -> Vec<OsString> {
    std::env::args_os()
        .map(|arg| {
            if arg == "-ns" {
                "--namespaces".into()
            } else {
                arg
            }
        })
        .collect()
}

fn open_source(input: &str) -> anyhow::Result<PageSource> {
    if input == "-" {
        Ok(PageSource::stdin())
    } else {
        PageSource::open(Path::new(input)).with_context(|| format!("cannot open {input}"))
    }
}

/// Extracts a single page by exact title to stdout via the title index.
///
/// Templates are, by default, expanded lazily through the same title index —
/// each template a page uses is fetched on demand, so no full template
/// database is built. `--no-templates` skips expansion; `--templates FILE`
/// forces the bulk database instead (built once and cached), for when an
/// exact full-dump-equivalent expansion is wanted.
fn lookup_title(
    input: &str,
    title: &str,
    templates_file: Option<&str>,
    no_templates: bool,
    config: &ExtractorConfig,
) -> anyhow::Result<()> {
    ensure!(input != "-", "--title needs a dump file on disk, not stdin");
    let dump_path = Path::new(input);

    let index = TitleIndex::open_or_build(dump_path)
        .with_context(|| format!("cannot open or build the title index for {input}"))?;
    let normalized = ruwex::title_index::normalize_title(title);
    if normalized != title {
        info!("normalized title to {normalized:?}");
    }
    let page = index
        .find_page(&normalized)?
        .with_context(|| format!("no page titled {normalized:?} found in {input}"))?;
    let site = index.site_info()?;

    // Choose the template source. `find_page`/`site_info` are already done, so
    // the index can be moved into the lazy source.
    let source: Box<dyn TemplateSource> = match (no_templates, templates_file) {
        (true, _) => Box::new(TemplateDb::default()),
        (false, Some(file)) => {
            let cache = Path::new(file);
            if !cache.is_file() {
                info!("building template cache {file} (first use; one-time full-dump scan)");
            }
            Box::new(ruwex::expand::templates::load(
                Some(dump_path),
                Some(cache),
                config.workers,
            )?)
        }
        (false, None) => Box::new(LazyTemplateSource::new(index)?),
    };

    let mut sink = StdoutSink::new();
    sink.write_doc(&render_page(&page, &site, config, source.as_ref()))?;
    sink.finish()?;
    Ok(())
}

fn main() {
    if let Err(error) = run() {
        eprintln!("error: {error:#}");
        std::process::exit(1);
    }
}

fn run() -> anyhow::Result<()> {
    let args = Args::parse_from(preprocessed_args());
    if args.version {
        println!("wikiextractor {} (ruwex)", env!("CARGO_PKG_VERSION"));
        return Ok(());
    }
    let input = args
        .input
        .context("the following arguments are required: input")?;

    let level = if args.debug {
        LevelFilter::Debug
    } else if args.quiet {
        LevelFilter::Warn
    } else {
        LevelFilter::Info
    };
    env_logger::Builder::new()
        .filter_level(level)
        .format_target(false)
        .init();

    let config = ExtractorConfig {
        format: if args.json {
            OutputFormat::Json
        } else {
            OutputFormat::Doc
        },
        // --html subsumes --links, as in the original
        keep_links: args.links || args.html,
        html: args.html,
        html_safe: args.html_safe.is_none_or(|v| {
            !matches!(
                v.trim().to_ascii_lowercase().as_str(),
                "" | "0" | "false" | "no"
            )
        }),
        namespaces: args
            .namespaces
            .map(|list| {
                list.split(',')
                    .map(|name| name.trim().to_string())
                    .filter(|name| !name.is_empty())
                    .collect()
            })
            .unwrap_or_else(ruwex::default_namespaces),
        workers: args.processes.unwrap_or_else(default_workers).max(1),
    };
    let max_bytes = parse_size(&args.bytes)?;

    if let Some(title) = args.title.as_deref() {
        return lookup_title(
            &input,
            title,
            args.templates.as_deref(),
            args.no_templates,
            &config,
        );
    }

    if args.article {
        // single-article analysis: output to stdout; templates only load
        // from an existing --templates file (like the original)
        let templates = match args
            .templates
            .as_deref()
            .map(Path::new)
            .filter(|p| p.is_file())
        {
            Some(cache) => Arc::new(ruwex::expand::templates::load(
                None,
                Some(cache),
                config.workers,
            )?),
            None => Arc::new(TemplateDb::default()),
        };
        let source = open_source(&input)?;
        let mut sink = StdoutSink::new();
        ruwex::run_with_templates(source, &templates, &config, &mut sink)?;
        return Ok(sink.finish()?);
    }

    let templates = if args.no_templates {
        Arc::new(TemplateDb::default())
    } else {
        let dump_path = (input != "-").then(|| Path::new(&input));
        let cache_path = args.templates.as_deref().map(Path::new);
        info!("preprocessing to collect template definitions; this may take some time");
        let load_started = Instant::now();
        let db = ruwex::expand::templates::load(dump_path, cache_path, config.workers)
            .context("cannot load template definitions")?;
        info!(
            "loaded {} templates in {:.1}s",
            db.len(),
            load_started.elapsed().as_secs_f64()
        );
        Arc::new(db)
    };

    let source = open_source(&input)?;

    let mut sink: Box<dyn DocSink> = if args.output == "-" {
        Box::new(StdoutSink::new())
    } else {
        Box::new(
            ShardedWriter::create(Path::new(&args.output), max_bytes, args.compress)
                .with_context(|| format!("cannot create output in {}", args.output))?,
        )
    };

    info!(
        "starting extraction of {} with {} workers",
        input, config.workers
    );
    let started = Instant::now();
    let stats = ruwex::run_with_templates(source, &templates, &config, sink.as_mut())?;
    sink.finish()?;

    let seconds = started.elapsed().as_secs_f64();
    info!(
        "finished: {} pages read, {} documents written in {:.1}s ({:.0} docs/s)",
        stats.pages,
        stats.docs,
        seconds,
        stats.docs as f64 / seconds.max(1e-9)
    );
    Ok(())
}