ruwex 0.1.0

//! The template database: definitions collected from `Template:` pages
//! (`define_template` port), plus loading — either a preprocessing pass over
//! the dump (parallel on multistream input) or a `--templates` cache file in
//! the same format Python wikiextractor writes, so cache files are
//! interchangeable between the two implementations.

use std::collections::{BTreeMap, HashMap, HashSet};
use std::fs::File;
use std::io::{BufRead, BufReader, BufWriter, Write};
use std::path::Path;
use std::sync::{Arc, RwLock};
use std::thread;

use bzip2::read::MultiBzDecoder;
use crossbeam_channel::bounded;
use regex::Regex;
use std::sync::LazyLock;

use super::functions::ucfirst;
use super::template::Template;
use crate::clean::tags::COMMENT;
use crate::dump::reader::PageSource;
use crate::dump::xml::DumpParser;
use crate::dump::{SiteInfo, multistream};
use crate::error::{Error, Result};
use crate::title_index::TitleIndex;

static REDIRECT: LazyLock<Regex> =
    LazyLock::new(|| Regex::new(r"(?i)^#REDIRECT.*?\[\[([^\]]*)\]\]").unwrap());
static NOINCLUDE: LazyLock<Regex> =
    LazyLock::new(|| Regex::new(r"(?s)<noinclude>(?:.*?)</noinclude>").unwrap());
static NOINCLUDE_UNTERMINATED: LazyLock<Regex> =
    LazyLock::new(|| Regex::new(r"(?s)<noinclude\s*>.*$").unwrap());
static ONLYINCLUDE: LazyLock<Regex> =
    LazyLock::new(|| Regex::new(r"(?s)<onlyinclude>(.*?)</onlyinclude>").unwrap());

pub struct TemplateDb {
    templates: HashMap<String, String>,
    redirects: HashMap<String, String>,
    /// Lazily parsed bodies, shared across workers.
    parsed: RwLock<HashMap<String, Arc<Template>>>,
    /// `Template:` (from siteinfo namespace 10).
    pub template_prefix: String,
    /// `Module:` (from siteinfo namespace 828); collected for cache files.
    pub module_prefix: String,
    /// All namespace names from siteinfo, used to qualify template titles.
    pub known_namespaces: HashSet<String>,
}

impl Default for TemplateDb {
    fn default() -> Self {
        Self {
            templates: HashMap::new(),
            redirects: HashMap::new(),
            parsed: RwLock::new(HashMap::new()),
            template_prefix: "Template:".to_string(),
            module_prefix: "Module:".to_string(),
            known_namespaces: HashSet::from(["Template".to_string()]),
        }
    }
}

/// Namespace facts derived from a dump's `<siteinfo>`, shared by both
/// template sources: the `Template:`/`Module:` prefixes and the set of known
/// namespace names used to qualify invocation titles.
pub(crate) struct Namespaces {
    pub template_prefix: String,
    pub module_prefix: String,
    pub known: HashSet<String>,
}

impl Namespaces {
    pub(crate) fn from_site_info(site: &SiteInfo) -> Self {
        let ns_name = |key: i32, default: &str| {
            site.namespaces
                .get(&key)
                .filter(|name| !name.is_empty())
                .cloned()
                .unwrap_or_else(|| default.to_string())
        };
        let mut known: HashSet<String> = site
            .namespaces
            .values()
            .filter(|n| !n.is_empty())
            .cloned()
            .collect();
        known.insert("Template".to_string());
        Self {
            template_prefix: format!("{}:", ns_name(10, "Template")),
            module_prefix: format!("{}:", ns_name(828, "Module")),
            known,
        }
    }
}

/// Port of `fullyQualifiedTemplateTitle`: resolves an invocation title to the
/// namespace of the page it includes. Shared by both template sources.
pub(crate) fn qualify_title(title: &str, known: &HashSet<String>, template_prefix: &str) -> String {
    if let Some(rest) = title.strip_prefix(':') {
        // a leading colon by itself implies the main namespace
        return ucfirst(rest);
    }
    if let Some(colon) = title.find(':') {
        let prefix = ucfirst(&title[..colon]);
        if known.contains(&prefix) {
            return format!("{prefix}{}", &title[colon..]);
        }
    }
    if title.is_empty() {
        String::new()
    } else {
        format!("{template_prefix}{}", ucfirst(title))
    }
}

impl TemplateDb {
    pub fn from_site_info(site: &SiteInfo) -> Self {
        let ns = Namespaces::from_site_info(site);
        Self {
            template_prefix: ns.template_prefix,
            module_prefix: ns.module_prefix,
            known_namespaces: ns.known,
            ..Default::default()
        }
    }

    pub fn len(&self) -> usize {
        self.templates.len()
    }

    pub fn is_empty(&self) -> bool {
        self.templates.is_empty()
    }

    /// Port of `define_template`: registers a `Template:` page, honoring
    /// redirects, comments, and noinclude/includeonly/onlyinclude sections.
    pub fn define(&mut self, title: &str, text: &str) {
        self.insert(title, prepare_definition(text));
    }

    fn insert(&mut self, title: &str, definition: Definition) {
        match definition {
            Definition::Redirect(target) => {
                self.redirects.insert(title.to_string(), target);
            }
            Definition::Body(text) => {
                if let Some(previous) = self.templates.get(title)
                    && *previous != text
                {
                    log::debug!("redefining template {title}");
                }
                self.templates.insert(title.to_string(), text);
            }
            Definition::Skip => {}
        }
    }

    /// Resolves a qualified title (through at most one redirect) to its
    /// parsed template, caching the parse.
    pub fn get_parsed(&self, qualified_title: &str) -> Option<Arc<Template>> {
        let title = self
            .redirects
            .get(qualified_title)
            .map(String::as_str)
            .unwrap_or(qualified_title);
        if let Some(template) = self.parsed.read().expect("lock poisoned").get(title) {
            return Some(Arc::clone(template));
        }
        let body = self.templates.get(title)?;
        let template = Arc::new(Template::parse(body));
        self.parsed
            .write()
            .expect("lock poisoned")
            .insert(title.to_string(), Arc::clone(&template));
        Some(template)
    }

    /// Port of `fullyQualifiedTemplateTitle`: resolves the namespace of a
    /// page included through the template mechanism.
    pub fn qualify(&self, template_title: &str) -> String {
        qualify_title(
            template_title,
            &self.known_namespaces,
            &self.template_prefix,
        )
    }
}

impl super::TemplateSource for TemplateDb {
    fn qualify(&self, title: &str) -> String {
        TemplateDb::qualify(self, title)
    }

    fn get_parsed(&self, qualified_title: &str) -> Option<Arc<Template>> {
        TemplateDb::get_parsed(self, qualified_title)
    }
}

/// A [`TemplateSource`](super::TemplateSource) that fetches template bodies on
/// demand through a [`TitleIndex`], instead of loading the whole template
/// database up front. Right for single-page lookups: it touches only the
/// templates a page actually uses, decompressing each containing block at
/// most once. Not intended for bulk extraction, where loading everything into
/// a [`TemplateDb`] avoids repeated seeks.
pub struct LazyTemplateSource {
    index: TitleIndex,
    template_prefix: String,
    known_namespaces: HashSet<String>,
    /// Decompressed stream blocks, `offset -> (title -> raw wikitext)`.
    blocks: RwLock<HashMap<u64, Arc<HashMap<String, String>>>>,
    /// Parsed templates by qualified title (`None` = known-absent), so a
    /// missing or repeated template is resolved at most once.
    parsed: RwLock<HashMap<String, Option<Arc<Template>>>>,
}

/// Redirect chains longer than this resolve to nothing (MediaWiki stops at
/// double redirects; a small bound also guards against redirect loops).
const MAX_REDIRECT_HOPS: usize = 3;

impl LazyTemplateSource {
    /// Wraps a title index (typically from `TitleIndex::open_or_build`),
    /// reading the dump's `<siteinfo>` once for its namespace prefixes.
    pub fn new(index: TitleIndex) -> Result<Self> {
        let ns = Namespaces::from_site_info(&index.site_info()?);
        Ok(Self {
            index,
            template_prefix: ns.template_prefix,
            known_namespaces: ns.known,
            blocks: RwLock::new(HashMap::new()),
            parsed: RwLock::new(HashMap::new()),
        })
    }

    /// Raw wikitext of a page by exact title, via the block cache.
    fn raw_body(&self, title: &str) -> Result<Option<String>> {
        let Some(offset) = self.index.offset(title) else {
            return Ok(None);
        };
        if let Some(block) = self.blocks.read().expect("lock poisoned").get(&offset) {
            return Ok(block.get(title).cloned());
        }
        let mut map = HashMap::new();
        for page in self.index.read_block(offset)? {
            map.insert(page.title, page.text);
        }
        let block = Arc::new(map);
        let body = block.get(title).cloned();
        self.blocks
            .write()
            .expect("lock poisoned")
            .insert(offset, block);
        Ok(body)
    }

    fn resolve(&self, title: &str, depth: usize) -> Option<Arc<Template>> {
        if let Some(cached) = self.parsed.read().expect("lock poisoned").get(title) {
            return cached.clone();
        }
        let result = self.fetch(title, depth);
        self.parsed
            .write()
            .expect("lock poisoned")
            .insert(title.to_string(), result.clone());
        result
    }

    fn fetch(&self, title: &str, depth: usize) -> Option<Arc<Template>> {
        if depth > MAX_REDIRECT_HOPS {
            return None;
        }
        let raw = self.raw_body(title).ok().flatten()?;
        match prepare_definition(&raw) {
            Definition::Redirect(target) => self.resolve(&target, depth + 1),
            Definition::Body(text) => Some(Arc::new(Template::parse(&text))),
            Definition::Skip => None,
        }
    }
}

impl super::TemplateSource for LazyTemplateSource {
    fn qualify(&self, title: &str) -> String {
        qualify_title(title, &self.known_namespaces, &self.template_prefix)
    }

    fn get_parsed(&self, qualified_title: &str) -> Option<Arc<Template>> {
        // Match the bulk database (and Python): only pages in the template
        // namespace are transcludable; `{{:Article}}` and the like resolve to
        // nothing rather than pulling in arbitrary content.
        if !qualified_title.starts_with(&self.template_prefix) {
            return None;
        }
        self.resolve(qualified_title, 0)
    }
}

/// Loads the template database for a dump.
///
/// - `dump_path`: the dump to preprocess; `None` when reading from stdin
///   (then `cache_file` must exist, as in the Python original).
/// - `cache_file`: if it exists it is loaded instead of scanning the dump;
///   if given but absent, the collected pages are saved to it.
pub fn load(
    dump_path: Option<&Path>,
    cache_file: Option<&Path>,
    workers: usize,
) -> Result<TemplateDb> {
    let cached = cache_file.filter(|f| f.is_file());
    let mut db = match dump_path {
        Some(path) => TemplateDb::from_site_info(&dump_site_info(path)?),
        None => {
            if cached.is_none() {
                return Err(Error::InvalidDump(
                    "to use templates with a stdin dump, supply an existing --templates file"
                        .to_string(),
                ));
            }
            TemplateDb::default()
        }
    };

    if let Some(cache) = cached {
        let pages = read_cache_pages(cache)?;
        log::info!(
            "loaded {} template/module pages from {}",
            pages.len(),
            cache.display()
        );
        apply_definitions(&mut db, &pages, workers);
        return Ok(db);
    }

    let path = dump_path.expect("checked above");
    let pages = collect_template_pages(path, &db, workers)?;
    if let Some(cache) = cache_file {
        save_cache(cache, &pages)?;
        log::info!(
            "saved {} template/module pages to {}",
            pages.len(),
            cache.display()
        );
    }
    apply_definitions(&mut db, &pages, workers);
    Ok(db)
}

/// The result of preprocessing a template page body for the database.
enum Definition {
    Redirect(String),
    Body(String),
    Skip,
}

/// The pure (parallelizable) part of `define_template`.
fn prepare_definition(text: &str) -> Definition {
    let first_line = text.lines().next().unwrap_or("");
    if let Some(caps) = REDIRECT.captures(first_line) {
        return Definition::Redirect(caps[1].to_string());
    }

    let text = COMMENT.replace_all(text, "");
    let text = NOINCLUDE.replace_all(&text, "");
    let text = NOINCLUDE_UNTERMINATED.replace_all(&text, "");
    let text = text.replace("<noinclude/>", "");

    // If <onlyinclude> parts are present, only they are transcluded;
    // otherwise the <includeonly> tags (not their content) are removed.
    let mut only = String::new();
    for caps in ONLYINCLUDE.captures_iter(&text) {
        only.push_str(&caps[1]);
    }
    let text = if only.is_empty() {
        text.replace("<includeonly>", "")
            .replace("</includeonly>", "")
    } else {
        only
    };

    if text.is_empty() {
        Definition::Skip
    } else {
        Definition::Body(text)
    }
}

/// Preprocesses page bodies in parallel (regex-heavy), then inserts them in
/// dump order so redefinition semantics stay deterministic.
fn apply_definitions(db: &mut TemplateDb, pages: &[(String, String)], workers: usize) {
    let template_pages: Vec<&(String, String)> = pages
        .iter()
        .filter(|(title, _)| title.starts_with(&db.template_prefix))
        .collect();
    let chunk_size = template_pages.len().div_ceil(workers.max(1)).max(1);
    let prepared: Vec<Vec<(&str, Definition)>> = thread::scope(|scope| {
        let handles: Vec<_> = template_pages
            .chunks(chunk_size)
            .map(|chunk| {
                scope.spawn(move || {
                    chunk
                        .iter()
                        .map(|(title, text)| (title.as_str(), prepare_definition(text)))
                        .collect()
                })
            })
            .collect();
        handles
            .into_iter()
            .map(|handle| handle.join().expect("definition worker panicked"))
            .collect()
    });
    for batch in prepared {
        for (title, definition) in batch {
            db.insert(title, definition);
        }
    }
}

/// Reads just the `<siteinfo>` header of a dump. Always reads sequentially:
/// the header sits at the start of the file, so this stays cheap even for
/// multistream dumps (and avoids parsing their large index files).
fn dump_site_info(path: &Path) -> Result<SiteInfo> {
    let file = File::open(path)?;
    let reader: Box<dyn BufRead> = if path.extension().is_some_and(|e| e == "bz2") {
        Box::new(BufReader::new(MultiBzDecoder::new(BufReader::new(file))))
    } else {
        Box::new(BufReader::new(file))
    };
    Ok(DumpParser::new(reader)
        .site_info()?
        .cloned()
        .unwrap_or_default())
}

/// Scans the dump for `Template:` and `Module:` pages (dump order preserved;
/// multistream input is scanned in parallel).
fn collect_template_pages(
    path: &Path,
    db: &TemplateDb,
    workers: usize,
) -> Result<Vec<(String, String)>> {
    let wanted = |title: &str| {
        title.starts_with(&db.template_prefix) || title.starts_with(&db.module_prefix)
    };
    match PageSource::open(path)? {
        PageSource::Sequential(reader) => {
            let mut parser = DumpParser::new(reader);
            let mut pages = Vec::new();
            while let Some(page) = parser.next_page()? {
                if wanted(&page.title) {
                    pages.push((page.title, page.text));
                }
            }
            Ok(pages)
        }
        PageSource::Multistream(ms) => {
            let workers = workers.max(1);
            let stream_count = ms.offsets.len() as u64;
            let (job_tx, job_rx) = bounded::<(u64, u64)>(workers * 4);
            let (result_tx, result_rx) =
                bounded::<(u64, Result<Vec<(String, String)>>)>(workers * 2);

            thread::scope(|scope| -> Result<Vec<(String, String)>> {
                let offsets = ms.offsets;
                scope.spawn(move || {
                    for (seq, offset) in offsets.into_iter().enumerate() {
                        if job_tx.send((seq as u64, offset)).is_err() {
                            break;
                        }
                    }
                });
                for _ in 0..workers {
                    let job_rx = job_rx.clone();
                    let result_tx = result_tx.clone();
                    let path = ms.path.clone();
                    scope.spawn(move || -> Result<()> {
                        let mut file = File::open(&path)?;
                        for (seq, offset) in job_rx {
                            let result = scan_stream(&mut file, offset, &wanted);
                            let failed = result.is_err();
                            if result_tx.send((seq, result)).is_err() || failed {
                                break;
                            }
                        }
                        Ok(())
                    });
                }
                drop(job_rx);
                drop(result_tx);

                let mut pending: BTreeMap<u64, Vec<(String, String)>> = BTreeMap::new();
                let mut pages = Vec::new();
                let mut next = 0;
                for (seq, result) in result_rx {
                    pending.insert(seq, result?);
                    while let Some(batch) = pending.remove(&next) {
                        pages.extend(batch);
                        next += 1;
                    }
                }
                if next != stream_count {
                    return Err(Error::InvalidDump(format!(
                        "template scan processed {next} of {stream_count} multistream blocks"
                    )));
                }
                Ok(pages)
            })
        }
    }
}

fn scan_stream(
    file: &mut File,
    offset: u64,
    wanted: &impl Fn(&str) -> bool,
) -> Result<Vec<(String, String)>> {
    let bytes = multistream::read_stream(file, offset)?;
    let mut parser = DumpParser::new(bytes.as_slice());
    let mut pages = Vec::new();
    while let Some(page) = parser.next_page()? {
        if wanted(&page.title) {
            pages.push((page.title, page.text));
        }
    }
    Ok(pages)
}

fn read_cache_pages(cache: &Path) -> Result<Vec<(String, String)>> {
    let file = File::open(cache)?;
    let reader: Box<dyn BufRead> = if cache.extension().is_some_and(|e| e == "bz2") {
        Box::new(BufReader::new(MultiBzDecoder::new(BufReader::new(file))))
    } else {
        Box::new(BufReader::new(file))
    };
    let mut parser = DumpParser::new(reader);
    let mut pages = Vec::new();
    while let Some(page) = parser.next_page()? {
        pages.push((page.title, page.text));
    }
    Ok(pages)
}

/// Writes collected pages in the same pseudo-dump format as the Python
/// original, with XML escaping so the file parses cleanly on reload.
fn save_cache(cache: &Path, pages: &[(String, String)]) -> Result<()> {
    let mut out = BufWriter::new(File::create(cache)?);
    for (title, text) in pages {
        writeln!(out, "<page>")?;
        writeln!(out, "   <title>{}</title>", xml_escape(title))?;
        writeln!(out, "   <ns>10</ns>")?;
        writeln!(out, "   <text>{}</text>", xml_escape(text))?;
        writeln!(out, "</page>")?;
    }
    out.flush()?;
    Ok(())
}

fn xml_escape(text: &str) -> String {
    text.replace('&', "&amp;")
        .replace('<', "&lt;")
        .replace('>', "&gt;")
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn define_handles_includeonly_and_noinclude() {
        let mut db = TemplateDb::default();
        db.define(
            "Template:T",
            "<noinclude>docs</noinclude><includeonly>body {{{1|}}}</includeonly>",
        );
        assert_eq!(db.templates["Template:T"], "body {{{1|}}}");

        db.define("Template:U", "keep<noinclude> dropped to end");
        assert_eq!(db.templates["Template:U"], "keep");
    }

    #[test]
    fn define_prefers_onlyinclude_sections() {
        let mut db = TemplateDb::default();
        db.define(
            "Template:O",
            "junk<onlyinclude>A</onlyinclude>mid<onlyinclude>B</onlyinclude>junk",
        );
        assert_eq!(db.templates["Template:O"], "AB");
    }

    #[test]
    fn define_records_redirects() {
        let mut db = TemplateDb::default();
        db.define(
            "Template:R",
            "#REDIRECT [[Template:Target]] {{R from move}}",
        );
        assert_eq!(db.redirects["Template:R"], "Template:Target");
        db.define("Template:Target", "real body");
        assert!(db.get_parsed("Template:R").is_some());
    }

    #[test]
    fn define_strips_comments() {
        let mut db = TemplateDb::default();
        db.define("Template:C", "a<!-- hidden -->b");
        assert_eq!(db.templates["Template:C"], "ab");
    }

    #[test]
    fn qualify_resolves_namespaces() {
        let db = TemplateDb::default();
        assert_eq!(db.qualify("lang-grc"), "Template:Lang-grc");
        assert_eq!(db.qualify("Template:X"), "Template:X");
        assert_eq!(db.qualify(":Main page"), "Main page");
        assert_eq!(db.qualify(""), "");
        // unknown namespace prefixes fall through to the template namespace
        assert_eq!(db.qualify("user:X"), "Template:User:X");
    }

    #[test]
    fn cache_roundtrip_preserves_definitions() {
        let dir = tempfile::tempdir().unwrap();
        let cache = dir.path().join("templates.cache");
        let pages = vec![(
            "Template:Esc".to_string(),
            "a < b & c > d\nline two".to_string(),
        )];
        save_cache(&cache, &pages).unwrap();

        let mut db = TemplateDb::default();
        let pages = read_cache_pages(&cache).unwrap();
        apply_definitions(&mut db, &pages, 2);
        assert_eq!(db.templates["Template:Esc"], "a < b & c > d\nline two");
    }
}