pdq 0.5.1

use std::{
    fs,
    path::Path,
    sync::{
        atomic::{AtomicBool, Ordering},
        Arc,
    },
};

use hayro::hayro_interpret::{
    font::Glyph,
    hayro_cmap::BfString,
    hayro_syntax::{
        object::{dict::keys::ANNOTS, Array},
        page::Page,
        LoadPdfError, Pdf,
    },
    interpret_page, BlendMode, ClipPath, Context, Device, GlyphDrawMode, Image, InterpreterCache,
    InterpreterSettings, InterpreterWarning, Paint, PathDrawMode, SoftMask, TransformExt,
};
use kurbo::{Affine, BezPath, Point, Rect, Vec2};

use crate::{
    hayro_stack::{guard_nesting_depth, with_deep_stack},
    range::{dedupe_preserving_order, PageRangeGroup},
    PdfOpsError, Result,
};

/// Emitted `y` is the approximate glyph top: baseline minus this fraction of
/// the on-page em size (pdf.js's text-layer fallback ascent).
const ASCENT_FACTOR: f64 = 0.8;
/// hayro glyph transforms take font units (thousandths of an em) as input.
const FONT_UNITS_PER_EM: f64 = 1000.0;
/// Annotations on one page past which their appearance streams are skipped.
///
/// hayro-syntax rebuilds an object stream's entire offset table on *every*
/// object it resolves out of it (`ObjectStream::new` in `xref.rs`), so
/// resolving a page's annotations costs annotations x objects-in-that-stream.
/// `corpus/pdfjs/bug1978317.pdf` puts 32,768 `/Link` annotations in a 13 MB
/// object stream and takes 25 s to yield no text at all, against 93 ms with
/// annotations off.
///
/// The busiest page in the 1,856-file corpus carries 171 annotations, so this
/// is a 24x margin over anything real; pages past it are link or popup spam,
/// which by construction has no `/AP` and so contributes no text. Skipping
/// marks the page [`PageText::degraded`] rather than dropping the fact
/// silently. It bounds the blowup, it does not remove it — the real fix is a
/// per-stream offset cache upstream (LaurenzV/hayro), and until that ships
/// `render` keeps the same pathology, where dropping annotations would change
/// pixels.
const MAX_ANNOTATIONS_PER_PAGE: usize = 4_096;

/// Gap past a glyph's advance (in em) that reads as a word space rather than
/// kerning or tracking, synthesized as ' ' for PDFs that encode word gaps as
/// TJ offsets instead of space glyphs (LaTeX). Poppler's `minWordBreakSpace`
/// and pdf.js's `SPACE_IN_FLOW_MIN_FACTOR` both sit near 0.1 em; kerning
/// pairs stay well below it, word spaces (~0.2-0.5 em) well above.
const WORD_GAP_MIN: f64 = 0.1;
/// Widest gap still treated as an in-flow word space (pdf.js's
/// `SPACE_IN_FLOW_MAX_FACTOR`); anything wider starts a new run.
const WORD_GAP_MAX: f64 = 0.6;

#[derive(Debug, Clone)]
pub struct ExtractTextOptions {
    pub pages: Option<PageRangeGroup>,
    pub password: Option<String>,
    /// Include the text of annotation and form-field appearance streams —
    /// a filled text widget's value, a FreeText note, a stamp — alongside the
    /// page content stream. Defaults to `true`, matching `pdftotext`,
    /// `mutool draw -F txt` and pdf-oxide. Annotations flagged Hidden
    /// (`/F` bit 2) are excluded either way.
    pub annotations: bool,
}

impl Default for ExtractTextOptions {
    fn default() -> Self {
        Self {
            pages: None,
            password: None,
            // The JSON is the whole output: there is no second layer for a
            // caller to consult, so dropping appearance streams silently
            // loses filled-in form values. `CopyOptions::copy_annotations`
            // defaults to true for the same reason.
            annotations: true,
        }
    }
}

/// A horizontal (or uniformly-oriented) sequence of glyphs positioned on the
/// page. Coordinates are PDF points at 72 dpi with a top-left origin, after
/// the same rotation and cropbox handling `render` applies. `x`/`y` is the
/// top-left of the run's axis-aligned bounding box; for horizontal text
/// `width` is the sum of glyph advances and `height` equals `font_size`,
/// while rotated/vertical text yields a narrow, tall box.
#[derive(Debug, Clone, PartialEq)]
pub struct TextRun {
    pub text: String,
    pub x: f32,
    pub y: f32,
    pub width: f32,
    pub height: f32,
    pub font_size: f32,
}

#[derive(Debug, Clone, PartialEq)]
pub struct PageText {
    pub page: usize,
    pub width: f32,
    pub height: f32,
    /// True when the page's text is known to be incomplete or unreliable: a
    /// glyph could not be mapped to Unicode (emitted as U+FFFD), the
    /// interpreter reported a font it could not decode, or the page carried
    /// more than `MAX_ANNOTATIONS_PER_PAGE` annotations and their appearance
    /// streams were skipped. Distinguishes "extraction failed" from "no text".
    pub degraded: bool,
    pub runs: Vec<TextRun>,
}

pub fn extract_text(input: &Path, options: &ExtractTextOptions) -> Result<Vec<PageText>> {
    with_deep_stack(|| extract_text_inner(input, options))
}

fn extract_text_inner(input: &Path, options: &ExtractTextOptions) -> Result<Vec<PageText>> {
    let data = fs::read(input)?;
    guard_nesting_depth(&data, input)?;
    let pdf = load_pdf(data, options.password.as_deref(), input)?;
    let pages = pdf.pages();

    let selected = match &options.pages {
        Some(range) => dedupe_preserving_order(&range.resolve(pages.len())?),
        None => (1..=pages.len()).collect(),
    };

    let cache = InterpreterCache::new();
    Ok(selected
        .iter()
        .map(|&page_number| {
            extract_page(
                &pages[page_number - 1],
                page_number,
                &cache,
                options.annotations,
            )
        })
        .collect())
}

fn load_pdf(data: Vec<u8>, password: Option<&str>, input: &Path) -> Result<Pdf> {
    let result = match password {
        Some(password) => Pdf::new_with_password(data, password),
        None => Pdf::new(data),
    };
    result.map_err(|err| match err {
        LoadPdfError::Decryption(_) => PdfOpsError::Password(match password {
            Some(_) => format!("invalid password for {}", input.display()),
            None => format!(
                "{} is encrypted; supply --password to decrypt it",
                input.display()
            ),
        }),
        LoadPdfError::Invalid => PdfOpsError::InvalidStructure(format!(
            "failed to parse PDF for text extraction: {}",
            input.display()
        )),
    })
}

fn extract_page<'a>(
    page: &Page<'a>,
    page_number: usize,
    cache: &InterpreterCache<'a>,
    annotations: bool,
) -> PageText {
    let (width, height) = page.render_dimensions();
    let initial_transform = page.initial_transform(true).to_kurbo();

    // Counting is cheap where resolving is not: `raw_iter` walks the array's
    // bytes and never dereferences the `n 0 R` entries, so this costs
    // microseconds even on the 32,768-annotation page it exists to catch.
    let annotation_flood = annotations
        && page.raw().get::<Array<'_>>(ANNOTS).is_some_and(|annots| {
            annots.raw_iter().take(MAX_ANNOTATIONS_PER_PAGE + 1).count() > MAX_ANNOTATIONS_PER_PAGE
        });
    let annotations = annotations && !annotation_flood;

    let font_warning = Arc::new(AtomicBool::new(false));
    let sink_flag = font_warning.clone();
    let settings = InterpreterSettings {
        warning_sink: Arc::new(move |warning| {
            if matches!(warning, InterpreterWarning::UnsupportedFont) {
                sink_flag.store(true, Ordering::Relaxed);
            }
        }),
        // Appearance streams carry text that appears nowhere else in the
        // output — filled form-field values above all — so they are part of
        // the text layer, as they are for pdftotext and mutool. hayro skips
        // annotations flagged Hidden (`/F` bit 2) but not NoView (bit 6),
        // and does not special-case `/Popup`; both are upstream gaps pdq
        // cannot filter from the outside (the annotation loop is inside
        // `interpret_page`).
        render_annotations: annotations,
        ..Default::default()
    };

    let mut context = Context::new(
        initial_transform,
        Rect::new(0.0, 0.0, width as f64, height as f64),
        cache,
        page.xref(),
        settings,
    );
    let mut device = TextDevice::default();
    interpret_page(page, &mut context, &mut device);

    let (runs, missing_unicode) = device.finish();
    PageText {
        page: page_number,
        width,
        height,
        degraded: missing_unicode || font_warning.load(Ordering::Relaxed) || annotation_flood,
        runs,
    }
}

struct RunState {
    text: String,
    /// Baseline origin of the first glyph.
    start: Point,
    /// Baseline point past the last glyph's advance.
    end: Point,
    font_size: f64,
    /// Unit vector pointing from the baseline toward the glyph top.
    up: Vec2,
    /// Unit vector along the advance direction.
    dir: Vec2,
    /// Where the next glyph's origin should land if the text continues
    /// uninterrupted; `None` when the glyph's advance width is unknown.
    expected: Option<Point>,
    last_origin: Point,
}

#[derive(Default)]
struct TextDevice {
    runs: Vec<TextRun>,
    current: Option<RunState>,
    missing_unicode: bool,
    last_fill_origin: Option<Point>,
}

impl TextDevice {
    fn finish(&mut self) -> (Vec<TextRun>, bool) {
        self.flush();
        (std::mem::take(&mut self.runs), self.missing_unicode)
    }

    fn flush(&mut self) {
        if let Some(run) = self.current.take() {
            if !run.text.trim().is_empty() {
                let ascent = run.up * (ASCENT_FACTOR * run.font_size);
                let descent = run.up * ((ASCENT_FACTOR - 1.0) * run.font_size);
                let corners = [
                    run.start + ascent,
                    run.start + descent,
                    run.end + ascent,
                    run.end + descent,
                ];
                let (mut min, mut max) = (corners[0], corners[0]);
                for corner in &corners[1..] {
                    min.x = min.x.min(corner.x);
                    min.y = min.y.min(corner.y);
                    max.x = max.x.max(corner.x);
                    max.y = max.y.max(corner.y);
                }
                self.runs.push(TextRun {
                    text: run.text,
                    x: min.x as f32,
                    y: min.y as f32,
                    width: (max.x - min.x) as f32,
                    height: (max.y - min.y) as f32,
                    font_size: run.font_size as f32,
                });
            }
        }
    }

    /// Whether the glyph at `origin` continues the current run; `Some(true)`
    /// means it continues after a word-sized gap (synthesize a space).
    fn continues_current(
        &self,
        origin: Point,
        font_size: f64,
        up: Vec2,
        dir: Vec2,
    ) -> Option<bool> {
        let current = self.current.as_ref()?;
        if (font_size - current.font_size).abs() > 0.02 * current.font_size
            || current.dir.dot(dir) < 0.99
            || current.up.dot(up) < 0.99
        {
            return None;
        }
        let (reference, gap_range, has_advance) = match current.expected {
            Some(expected) => (expected, (-0.25, WORD_GAP_MAX), true),
            // Without advance metrics (Type3 fonts), measure from the
            // previous origin and accept anything up to a wide glyph plus
            // a word gap; the glyph width is indistinguishable from the
            // gap, so no space can be synthesized.
            None => (current.last_origin, (-0.05, 1.6), false),
        };
        let delta = origin - reference;
        let along = delta.dot(current.dir) / current.font_size;
        let perpendicular = delta.dot(current.up) / current.font_size;
        if perpendicular.abs() > 0.15 || along < gap_range.0 || along > gap_range.1 {
            return None;
        }
        Some(has_advance && along >= WORD_GAP_MIN)
    }
}

impl<'a> Device<'a> for TextDevice {
    fn draw_glyph(
        &mut self,
        glyph: &Glyph<'a>,
        transform: Affine,
        glyph_transform: Affine,
        _paint: &Paint<'a>,
        draw_mode: &GlyphDrawMode,
    ) {
        let to_device = transform * glyph_transform;
        let origin = to_device * Point::ZERO;
        let em_x = to_device * Point::new(FONT_UNITS_PER_EM, 0.0) - origin;
        let em_y = to_device * Point::new(0.0, FONT_UNITS_PER_EM) - origin;
        let font_size = em_y.hypot();
        if !origin.x.is_finite()
            || !origin.y.is_finite()
            || !font_size.is_finite()
            || font_size <= 0.0
            || em_x.hypot() <= 0.0
        {
            return;
        }

        // FillStroke rendering modes draw the same glyph twice; keep only
        // the fill pass. Stroke-only text still comes through because no
        // fill preceded it at the same origin.
        if matches!(draw_mode, GlyphDrawMode::Stroke(_)) {
            if self.last_fill_origin == Some(origin) {
                return;
            }
        } else {
            self.last_fill_origin = Some(origin);
        }

        let text = match glyph.as_unicode() {
            Some(BfString::Char(c)) => c.to_string(),
            Some(BfString::String(s)) => s,
            None => {
                self.missing_unicode = true;
                '\u{FFFD}'.to_string()
            }
        };

        let up = em_y / font_size;
        let dir = em_x.normalize();
        let advance = match glyph {
            Glyph::Outline(outline) => outline
                .advance_width()
                .map(|w| to_device * Point::new(w as f64, 0.0) - origin),
            Glyph::Type3(_) => None,
        };
        let expected = advance.map(|a| origin + a);
        // Without an advance, estimate the glyph's extent as half an em so
        // the box still covers most of the final glyph.
        let end = expected.unwrap_or(origin + dir * (0.5 * font_size));

        if let Some(word_gap) = self.continues_current(origin, font_size, up, dir) {
            let current = self.current.as_mut().unwrap();
            if word_gap
                && !current.text.ends_with(char::is_whitespace)
                && !text.starts_with(char::is_whitespace)
            {
                current.text.push(' ');
            }
            current.text.push_str(&text);
            current.last_origin = origin;
            current.expected = expected;
            current.end = end;
        } else {
            self.flush();
            self.current = Some(RunState {
                text,
                start: origin,
                end,
                font_size,
                up,
                dir,
                expected,
                last_origin: origin,
            });
        }
    }

    fn set_soft_mask(&mut self, _: Option<SoftMask<'a>>) {}
    fn set_blend_mode(&mut self, _: BlendMode) {}
    fn draw_path(&mut self, _: &BezPath, _: Affine, _: &Paint<'a>, _: &PathDrawMode) {}
    fn push_clip_path(&mut self, _: &ClipPath) {}
    fn push_transparency_group(&mut self, _: f32, _: Option<SoftMask<'a>>, _: BlendMode) {}
    fn draw_image(&mut self, _: Image<'a, '_>, _: Affine) {}
    fn pop_clip_path(&mut self) {}
    fn pop_transparency_group(&mut self) {}
}

/// Serialize pages as the `pdq text` stdout JSON (an array of page objects).
pub fn pages_to_json(pages: &[PageText]) -> String {
    let mut out = String::from("[");
    for (i, page) in pages.iter().enumerate() {
        if i > 0 {
            out.push(',');
        }
        out.push_str(&format!(
            "{{\"page\":{},\"page_width\":{},\"page_height\":{},\"degraded\":{},\"runs\":[",
            page.page,
            format_number(page.width),
            format_number(page.height),
            page.degraded
        ));
        for (j, run) in page.runs.iter().enumerate() {
            if j > 0 {
                out.push(',');
            }
            out.push_str(&format!(
                "{{\"text\":\"{}\",\"x\":{},\"y\":{},\"width\":{},\"height\":{},\"font_size\":{}}}",
                escape_json(&run.text),
                format_number(run.x),
                format_number(run.y),
                format_number(run.width),
                format_number(run.height),
                format_number(run.font_size)
            ));
        }
        out.push_str("]}");
    }
    out.push(']');
    out
}

/// Numbers rounded to 1/1000 pt; f64 Display keeps them free of float noise.
fn format_number(value: f32) -> String {
    let rounded = (value as f64 * 1000.0).round() / 1000.0;
    if rounded.is_finite() {
        format!("{rounded}")
    } else {
        "0".to_string()
    }
}

fn escape_json(text: &str) -> String {
    let mut out = String::with_capacity(text.len());
    for c in text.chars() {
        match c {
            '"' => out.push_str("\\\""),
            '\\' => out.push_str("\\\\"),
            '\n' => out.push_str("\\n"),
            '\r' => out.push_str("\\r"),
            '\t' => out.push_str("\\t"),
            c if (c as u32) < 0x20 => out.push_str(&format!("\\u{:04x}", c as u32)),
            c => out.push(c),
        }
    }
    out
}

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

    #[test]
    fn json_escapes_control_and_quote_characters() {
        assert_eq!(escape_json("a\"b\\c\nd\u{1}e"), "a\\\"b\\\\c\\nd\\u0001e");
    }

    #[test]
    fn numbers_render_without_float_noise() {
        assert_eq!(format_number(72.0), "72");
        assert_eq!(format_number(57.599_996), "57.6");
        assert_eq!(format_number(f32::NAN), "0");
    }

    #[test]
    fn json_shape_matches_schema() {
        let pages = vec![PageText {
            page: 1,
            width: 612.0,
            height: 792.0,
            degraded: false,
            runs: vec![TextRun {
                text: "Invoice".to_string(),
                x: 72.0,
                y: 70.5,
                width: 57.0,
                height: 18.0,
                font_size: 18.0,
            }],
        }];
        assert_eq!(
            pages_to_json(&pages),
            "[{\"page\":1,\"page_width\":612,\"page_height\":792,\"degraded\":false,\
             \"runs\":[{\"text\":\"Invoice\",\"x\":72,\"y\":70.5,\
             \"width\":57,\"height\":18,\"font_size\":18}]}]"
        );
    }
}