formal-ai 0.92.0

//! Translation pipeline.
//!
//! Public entry point: [`TranslationPipeline::translate`].
//!
//! ## Stages
//!
//! 1. **Formalize** — `formalize(source_surface, source_lang)` projects the
//!    natural-language surface into a [`MeaningId`] by consulting
//!    Wiktionary (for the page that defines the lemma) and Wikidata
//!    (for the language-neutral Q-item or Lexeme sense). The
//!    [`MeaningId`] is the semantic meta-language identifier; it never
//!    embeds any single language.
//!
//! 2. **Deformalize** — `deformalize(meaning_id, target_lang)` renders a
//!    surface form in the target language by:
//!
//!    a. running a Wikidata SPARQL lexeme-join (P5137) when the meaning
//!    is backed by a Wikidata lexeme, **or**
//!
//!    b. parsing the source-edition Wiktionary translation table for
//!    the target language code, **or**
//!
//!    c. parsing the target-edition Wiktionary's `=== Перевод ===` /
//!    `===Translations===` block in reverse.
//!
//! Each strategy is recorded as a `provenance` entry on
//! [`Translation`], so the resulting links-notation trace shows
//! exactly which API responses fed the answer.

use super::http::{HttpClient, HttpError};
use super::meaning::MeaningId;
use super::wikidata::Wikidata;
use super::wiktionary::{Wiktionary, WiktionaryCandidate};

/// Result of a translation request.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Translation {
    pub source_surface: String,
    pub source_lang: String,
    pub target_lang: String,
    pub meaning: MeaningId,
    pub candidates: Vec<WiktionaryCandidate>,
    pub provenance: Vec<String>,
}

impl Translation {
    /// The best target-language surface form. Picks the first
    /// unqualified candidate, falling back to the first candidate.
    #[must_use]
    pub fn primary_surface(&self) -> Option<&str> {
        self.candidates
            .iter()
            .find(|c| c.qualifier.is_none())
            .or_else(|| self.candidates.first())
            .map(|c| c.surface.as_str())
    }
}

/// Translation pipeline. Borrows a single HTTP client (typically a
/// `CachedHttpClient`) for every Wiktionary and Wikidata call.
pub struct TranslationPipeline<'a, T: HttpClient + ?Sized> {
    http: &'a T,
}

impl<'a, T: HttpClient + ?Sized> TranslationPipeline<'a, T> {
    pub const fn new(http: &'a T) -> Self {
        Self { http }
    }

    /// Translate `surface` from `source_lang` to `target_lang`.
    ///
    /// Returns `Ok(Translation)` even when the candidate list is empty
    /// — callers can detect translation gaps by inspecting
    /// [`Translation::candidates`].
    pub fn translate(
        &self,
        surface: &str,
        source_lang: &str,
        target_lang: &str,
    ) -> Result<Translation, HttpError> {
        if source_lang == target_lang {
            return Ok(Translation {
                source_surface: surface.to_owned(),
                source_lang: source_lang.to_owned(),
                target_lang: target_lang.to_owned(),
                meaning: MeaningId::from_wiktionary_page(source_lang, surface),
                candidates: vec![WiktionaryCandidate {
                    surface: surface.to_owned(),
                    qualifier: None,
                }],
                provenance: vec!["identity".to_owned()],
            });
        }

        let mut provenance: Vec<String> = Vec::new();
        let page_title = normalize_page_title(surface);

        // Stage 1: formalize — fetch the source-edition Wiktionary page.
        // This single fetch usually carries translations for every target
        // language, grouped by sense (one block per `{{trans-top}}` on
        // English Wiktionary, one per `{{перев-блок}}` on Russian).
        let source_wiktionary = Wiktionary::new(source_lang, self.http);
        let mut blocks: Vec<Vec<WiktionaryCandidate>> = Vec::new();
        let mut meaning = MeaningId::from_wiktionary_page(source_lang, &page_title);
        match source_wiktionary.translation_blocks(&page_title, target_lang) {
            Ok(found) => {
                provenance.push(format!(
                    "wiktionary:{source_lang}:{page_title}#translations->{target_lang}"
                ));
                if !found.is_empty() {
                    blocks = found;
                }
            }
            Err(error) => {
                provenance.push(format!(
                    "wiktionary:{source_lang}:{page_title}#translations->error({error})",
                ));
            }
        }

        // Stage 1a: many high-traffic pages keep translations on a
        // `/translations` subpage (`{{see translation subpage|...}}`).
        // Try that subpage when the main page yielded nothing.
        if blocks.is_empty() {
            let subpage = format!("{page_title}/translations");
            match source_wiktionary.translation_blocks(&subpage, target_lang) {
                Ok(found) if !found.is_empty() => {
                    provenance.push(format!(
                        "wiktionary:{source_lang}:{subpage}#translations->{target_lang}"
                    ));
                    blocks = found;
                }
                Ok(_) => {}
                Err(error) => {
                    provenance.push(format!(
                        "wiktionary:{source_lang}:{subpage}#translations->error({error})",
                    ));
                }
            }
        }

        // Stage 1b: when the source-edition table came up empty, try the
        // target-edition Wiktionary in reverse. The target page lists
        // the source language under its own translation block — useful
        // when the source edition is sparse (common for ru → en).
        if blocks.is_empty() {
            if let Some(reverse) = reverse_lookup(
                self.http,
                surface,
                source_lang,
                target_lang,
                &mut provenance,
            ) {
                blocks = reverse;
            }
        }

        // Stage 1d: phrasal-variant fallback. Some natural phrases don't
        // have their own Wiktionary entry (e.g. ru:"как у тебя дела") but
        // a shorter canonical form does (ru:"как дела"). Generate
        // alternate forms by stripping language-specific filler
        // sub-phrases and retry Stages 1 + 1a + 1b on each variant. Stop
        // at the first variant that produces at least one block.
        let mut active_page_title = page_title.clone();
        if blocks.is_empty() {
            for variant in phrasal_variants(&page_title, source_lang) {
                provenance.push(format!("wiktionary:{source_lang}:variant->{variant}"));
                match source_wiktionary.translation_blocks(&variant, target_lang) {
                    Ok(found) if !found.is_empty() => {
                        provenance.push(format!(
                            "wiktionary:{source_lang}:{variant}#translations->{target_lang}"
                        ));
                        blocks = found;
                        active_page_title.clone_from(&variant);
                        meaning = MeaningId::from_wiktionary_page(source_lang, &variant);
                        break;
                    }
                    Ok(_) => {}
                    Err(error) => {
                        provenance.push(format!(
                            "wiktionary:{source_lang}:{variant}#translations->error({error})"
                        ));
                    }
                }
                if let Some(reverse) = reverse_lookup(
                    self.http,
                    &variant,
                    source_lang,
                    target_lang,
                    &mut provenance,
                ) {
                    blocks = reverse;
                    active_page_title.clone_from(&variant);
                    meaning = MeaningId::from_wiktionary_page(source_lang, &variant);
                    break;
                }
            }
        }

        // Stage 1c: pick the most likely sense block when the source word
        // is polysemous (multiple blocks). For each block, count how
        // many candidates round-trip — that is, how many target-edition
        // pages list the source surface as a translation. The block with
        // the highest confirmation rate wins; ties fall back to source
        // order.
        let candidates = if blocks.is_empty() {
            Vec::new()
        } else {
            select_best_block(
                self.http,
                &active_page_title,
                source_lang,
                target_lang,
                &mut provenance,
                blocks,
            )
        };

        // Stage 2: pull the Wikidata Q-item / Lexeme sense — even if
        // we already have a translation surface, we want the trace to
        // expose the canonical meaning id. Failures here do not block
        // the translation.
        let mut candidates = candidates;
        if let Some(updated) = upgrade_meaning_via_wikidata(
            self.http,
            &active_page_title,
            source_lang,
            target_lang,
            &mut provenance,
            &mut candidates,
        ) {
            meaning = updated;
        }

        Ok(Translation {
            source_surface: surface.to_owned(),
            source_lang: source_lang.to_owned(),
            target_lang: target_lang.to_owned(),
            meaning,
            candidates,
            provenance,
        })
    }
}

/// Try the target-edition Wiktionary in reverse. We open the page that
/// is most likely to *exist* on the target edition for the source word.
/// Russian Wiktionary, for instance, hosts pages for Russian words and
/// records their translations into other languages under
/// `=== Перевод ===`.
fn reverse_lookup<T: HttpClient + ?Sized>(
    http: &T,
    surface: &str,
    source_lang: &str,
    target_lang: &str,
    provenance: &mut Vec<String>,
) -> Option<Vec<Vec<WiktionaryCandidate>>> {
    let page_title = normalize_page_title(surface);
    for edition in [source_lang, target_lang] {
        let wiktionary = Wiktionary::new(edition, http);
        match wiktionary.translation_blocks(&page_title, target_lang) {
            Ok(blocks) if !blocks.is_empty() => {
                provenance.push(format!(
                    "wiktionary:{edition}:{page_title}#reverse->{target_lang}"
                ));
                return Some(blocks);
            }
            Ok(_) => {}
            Err(error) => {
                provenance.push(format!(
                    "wiktionary:{edition}:{page_title}#reverse->error({error})"
                ));
            }
        }
    }
    None
}

/// Pick the best sense block from `blocks` by round-trip confirmation rate.
///
/// Each candidate's "round-trip position" is its index within the
/// target-edition block that contains the source surface (not the
/// flattened candidate list). Lower position = the target considers the
/// source more primary.
///
/// Block-level selection picks the block with the most round-trip
/// confirmations (ties broken by source order). Within the chosen block,
/// candidates are sorted by `(target_position_within_block, source_idx)`.
///
/// Worked example — `en:hello`:
/// - Block 0 ("greeting"): `привет` confirms at target-pos 0,
///   `здравствуйте` at target-pos 1. Two confirms.
/// - Block 1 ("when answering the telephone"): `алло` and `алё` confirm
///   at target-pos 0 and 1. Two confirms.
/// - Ties go to the earlier source block → block 0. Within block 0,
///   `привет` (target-pos 0) beats `здравствуйте` (target-pos 1).
fn select_best_block<T: HttpClient + ?Sized>(
    http: &T,
    page_title: &str,
    source_lang: &str,
    target_lang: &str,
    provenance: &mut Vec<String>,
    blocks: Vec<Vec<WiktionaryCandidate>>,
) -> Vec<WiktionaryCandidate> {
    let target_wiktionary = Wiktionary::new(target_lang, http);
    let mut block_positions: Vec<Vec<Option<usize>>> = Vec::with_capacity(blocks.len());
    for (block_idx, block) in blocks.iter().enumerate() {
        let mut positions: Vec<Option<usize>> = Vec::with_capacity(block.len());
        for candidate in block {
            let candidate_page = normalize_page_title(&candidate.surface);
            if candidate_page.is_empty() {
                positions.push(None);
                continue;
            }
            let Ok(back_blocks) =
                target_wiktionary.translation_blocks(&candidate_page, source_lang)
            else {
                positions.push(None);
                continue;
            };
            // Within-block position: scan each block of the target page
            // and record the position of the source surface inside the
            // first block that contains it.
            let mut within_block_position: Option<usize> = None;
            for back_block in &back_blocks {
                if let Some(pos) = back_block
                    .iter()
                    .position(|row| normalize_page_title(&row.surface) == page_title)
                {
                    within_block_position = Some(pos);
                    break;
                }
            }
            if let Some(pos) = within_block_position {
                provenance.push(format!(
                    "wiktionary:{target_lang}:{candidate_page}#confirms->{source_lang}:{page_title}@{pos}[block{block_idx}]"
                ));
            }
            positions.push(within_block_position);
        }
        block_positions.push(positions);
    }
    let mut best_block: usize = 0;
    let mut best_confirms: usize = 0;
    for (idx, positions) in block_positions.iter().enumerate() {
        let confirms = positions.iter().filter(|p| p.is_some()).count();
        if confirms > best_confirms {
            best_confirms = confirms;
            best_block = idx;
        }
    }
    let block = blocks.into_iter().nth(best_block).unwrap_or_default();
    let positions = block_positions
        .into_iter()
        .nth(best_block)
        .unwrap_or_default();
    let mut indexed: Vec<(usize, Option<usize>, WiktionaryCandidate)> = block
        .into_iter()
        .zip(positions)
        .enumerate()
        .map(|(idx, (cand, pos))| (idx, pos, cand))
        .collect();
    indexed.sort_by_key(|(idx, pos, _)| {
        pos.as_ref()
            .map_or((1usize, 0, *idx), |p| (0usize, *p, *idx))
    });
    indexed.into_iter().map(|(_, _, cand)| cand).collect()
}

/// Upgrade a Wiktionary-only [`MeaningId`] to a Wikidata-backed one when
/// SPARQL returns a lexeme that matches the surface in the source
/// language. Also appends any extra lemma forms returned by the
/// translation query — they are stored as additional candidates so
/// callers can disambiguate (formal vs informal, dialect, etc).
fn upgrade_meaning_via_wikidata<T: HttpClient + ?Sized>(
    http: &T,
    page_title: &str,
    source_lang: &str,
    target_lang: &str,
    provenance: &mut Vec<String>,
    candidates: &mut Vec<WiktionaryCandidate>,
) -> Option<MeaningId> {
    let wikidata = Wikidata::new(http);
    let hits = match wikidata.search_lexeme(page_title, source_lang) {
        Ok(hits) => hits,
        Err(error) => {
            provenance.push(format!("wikidata:search->error({error})"));
            return None;
        }
    };
    let first = hits.first()?;
    provenance.push(format!("wikidata:lexeme:{}", first.id));
    let lemmas = match wikidata.lexeme_translations(&first.id, target_lang) {
        Ok(rows) => rows,
        Err(error) => {
            provenance.push(format!("wikidata:sparql->error({error})"));
            return Some(MeaningId::from_sense(first.id.clone()));
        }
    };
    if !lemmas.is_empty() {
        provenance.push(format!(
            "wikidata:sparql:{}->{} ({} lemmas)",
            first.id,
            target_lang,
            lemmas.len()
        ));
    }
    for lemma in lemmas {
        let candidate = WiktionaryCandidate {
            surface: lemma.value,
            qualifier: None,
        };
        if !candidates.iter().any(|c| c.surface == candidate.surface) {
            candidates.push(candidate);
        }
    }
    Some(MeaningId::from_sense(first.id.clone()))
}

/// Generate alternate phrasal forms to retry when the original page is
/// missing on Wiktionary.
///
/// Languages differ in which optional words can be elided without
/// changing meaning. We encode the elisions per source language; for now
/// we cover Russian (`у тебя/у вас/у нас/...` infix between question word
/// and noun).
///
/// Variants are returned in priority order — the caller stops at the
/// first one that produces translations.
#[must_use]
pub fn phrasal_variants(page_title: &str, source_lang: &str) -> Vec<String> {
    let mut variants: Vec<String> = Vec::new();
    if source_lang.eq_ignore_ascii_case("ru") {
        // Russian dative-of-possession infix: `как у тебя дела` and
        // `как у вас дела` collapse to the canonical `как дела`. This
        // covers the politeness variants without changing semantics.
        let pronouns = [
            "у тебя",
            "у вас",
            "у нас",
            "у меня",
            "у них",
            "у него",
            "у неё",
            "у нее",
        ];
        for pronoun in &pronouns {
            // Strip the pronoun as a whole-word infix surrounded by spaces.
            let needle = format!(" {pronoun} ");
            if let Some(idx) = page_title.find(&needle) {
                let mut alt = String::with_capacity(page_title.len() - needle.len() + 1);
                alt.push_str(&page_title[..idx]);
                alt.push(' ');
                alt.push_str(&page_title[idx + needle.len()..]);
                let alt = alt.split_whitespace().collect::<Vec<_>>().join(" ");
                if !alt.is_empty() && alt != page_title && !variants.contains(&alt) {
                    variants.push(alt);
                }
            }
        }
    }
    variants
}

/// Normalize a surface fragment into a Wiktionary page title.
///
/// - Trim whitespace.
/// - Strip terminal punctuation (`?`, `!`, `.`, `。`, `？`, `！`).
/// - Lower-case the first letter (Wiktionary stores most page titles in
///   lower case for content words; redirects handle the capitalized
///   variants but we want a stable cache key).
#[must_use]
pub fn normalize_page_title(surface: &str) -> String {
    let trimmed = surface
        .trim()
        .trim_end_matches(['?', '!', '.', '。', '？', '！', '．']);
    let mut chars = trimmed.chars();
    let Some(first) = chars.next() else {
        return String::new();
    };
    let mut out = String::with_capacity(trimmed.len());
    for character in first.to_lowercase() {
        out.push(character);
    }
    out.extend(chars);
    out
}

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

    struct StubHttp {
        responses: Mutex<std::collections::HashMap<String, String>>,
    }

    impl StubHttp {
        fn new(pairs: &[(&str, &str)]) -> Self {
            Self {
                responses: Mutex::new(
                    pairs
                        .iter()
                        .map(|(k, v)| ((*k).to_owned(), (*v).to_owned()))
                        .collect(),
                ),
            }
        }
    }

    impl HttpClient for StubHttp {
        fn get(&self, url: &str) -> Result<String, HttpError> {
            self.responses
                .lock()
                .unwrap()
                .get(url)
                .cloned()
                .ok_or_else(|| HttpError::Status {
                    status: 404,
                    body: format!("no stubbed response for {url}"),
                })
        }
    }

    #[test]
    fn normalize_page_title_strips_terminal_punctuation() {
        assert_eq!(normalize_page_title("Hello!"), "hello");
        assert_eq!(normalize_page_title("как у тебя дела?"), "как у тебя дела");
        assert_eq!(normalize_page_title("你好？"), "你好");
    }

    #[test]
    fn normalize_page_title_lowercases_first_letter() {
        assert_eq!(normalize_page_title("Hello"), "hello");
        assert_eq!(normalize_page_title("Как дела"), "как дела");
    }

    #[test]
    fn translate_identity_returns_self_with_identity_provenance() {
        let http = StubHttp::new(&[]);
        let pipeline = TranslationPipeline::new(&http);
        let translation = pipeline.translate("hello", "en", "en").unwrap();
        assert_eq!(translation.primary_surface(), Some("hello"));
        assert_eq!(translation.provenance, vec!["identity".to_owned()]);
    }

    #[test]
    fn translate_uses_source_edition_translation_table() {
        // English Wiktionary returns a JSON envelope around wikitext;
        // the wikitext lists the Russian translation under `{{t+|ru|...}}`.
        let url = "https://en.wiktionary.org/w/api.php?action=parse&page=hello&prop=wikitext&formatversion=2&format=json&redirects=1";
        let wikitext = r#"{"parse":{"title":"hello","wikitext":"* Russian: {{t+|ru|привет}}\n"}}"#;
        let http = StubHttp::new(&[(url, wikitext)]);
        let pipeline = TranslationPipeline::new(&http);
        let translation = pipeline.translate("hello", "en", "ru").unwrap();
        assert_eq!(translation.primary_surface(), Some("привет"));
        assert!(
            translation
                .provenance
                .iter()
                .any(|p| p.starts_with("wiktionary:en:hello#translations->ru")),
            "got provenance: {:?}",
            translation.provenance,
        );
    }

    #[test]
    fn translate_returns_translation_with_empty_candidates_when_nothing_matches() {
        // No HTTP stubs => every fetch fails. The pipeline still
        // produces a Translation, but with an empty candidates list
        // (callers detect the gap explicitly).
        let http = StubHttp::new(&[]);
        let pipeline = TranslationPipeline::new(&http);
        let translation = pipeline.translate("xyzzy", "en", "ru").unwrap();
        assert!(translation.candidates.is_empty());
        assert!(translation.primary_surface().is_none());
        assert!(translation.provenance.iter().any(|p| p.contains("error")));
    }

    #[test]
    fn translate_prefers_unqualified_candidate() {
        let url = "https://en.wiktionary.org/w/api.php?action=parse&page=hello&prop=wikitext&formatversion=2&format=json&redirects=1";
        let wikitext = r#"{"parse":{"wikitext":"* Russian: {{t|ru|здравствуйте|q=formal}}, {{t+|ru|привет|q=informal}}, {{t|ru|здорово}}\n"}}"#;
        let http = StubHttp::new(&[(url, wikitext)]);
        let pipeline = TranslationPipeline::new(&http);
        let translation = pipeline.translate("hello", "en", "ru").unwrap();
        // The first unqualified candidate wins.
        assert_eq!(translation.primary_surface(), Some("здорово"));
    }
}