inkhaven 2.4.0

//! 1.3.0 PDF-1 P3 — watermark / stamp (RFC §8.7).
//!
//! Stamp text (`DRAFT`, a name, a date) and/or an image (a logo) onto a
//! page range.  Each stamp is appended as a self-contained `q … Q` block
//! after the page's existing content — wrapped in its own graphics state
//! with a constant-alpha `ExtGState` for translucency — so it never
//! disturbs the body.  Text is centred and rotatable; the image is
//! centred and scaled to a fraction of the page width.

use std::path::PathBuf;

use lopdf::{Dictionary, Object};

use super::doc::PdfDoc;
use super::geometry::mm_to_pt;
use super::ops::PageSpec;
use super::{Error, Result};

/// Where on the page the stamp anchors.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum WmPosition {
    Center,
    TopLeft,
    TopRight,
    BottomLeft,
    BottomRight,
}

impl WmPosition {
    pub fn parse(s: &str) -> Option<Self> {
        match s.trim().to_ascii_lowercase().replace(['-', '_'], "").as_str() {
            "center" | "centre" | "middle" => Some(Self::Center),
            "topleft" | "tl" => Some(Self::TopLeft),
            "topright" | "tr" => Some(Self::TopRight),
            "bottomleft" | "bl" => Some(Self::BottomLeft),
            "bottomright" | "br" => Some(Self::BottomRight),
            _ => None,
        }
    }
}

/// What to stamp.  At least one of `text` / `image` should be set.
#[derive(Debug, Clone)]
pub struct WatermarkSpec {
    pub text: Option<String>,
    pub image: Option<PathBuf>,
    /// Constant alpha 0..=1 (fill + stroke).
    pub opacity: f32,
    /// Text rotation in degrees (counter-clockwise; e.g. 45 for a
    /// diagonal `DRAFT`).
    pub rotation_deg: f32,
    pub font_size_pt: f32,
    /// Text colour (0..=1 RGB).
    pub color: (f32, f32, f32),
    pub position: WmPosition,
    /// Image width as a fraction of the page width (aspect preserved).
    pub image_scale: f32,
    /// Page range; `None` = every page.
    pub pages: Option<PageSpec>,
}

impl Default for WatermarkSpec {
    fn default() -> Self {
        Self {
            text: None,
            image: None,
            opacity: 0.18,
            rotation_deg: 45.0,
            font_size_pt: 72.0,
            color: (0.5, 0.5, 0.5),
            position: WmPosition::Center,
            image_scale: 0.5,
            pages: None,
        }
    }
}

/// Inset from the page edge for the corner anchors.
const MARGIN: f32 = 24.0; // pt (~8.5 mm)

/// Stamp `spec` onto the selected pages.  Returns the number stamped.
pub fn apply_watermark(doc: &mut PdfDoc, spec: &WatermarkSpec) -> Result<usize> {
    if spec.text.is_none() && spec.image.is_none() {
        return Err(Error::Other("watermark: nothing to stamp (no text or image)".into()));
    }
    let page_ids = doc.page_ids().to_vec();
    let count = page_ids.len();
    let selected: Vec<u32> = spec
        .pages
        .as_ref()
        .map(|s| s.resolve(count))
        .unwrap_or_else(|| (1..=count as u32).collect());

    // Unicode text path: if the watermark text is non-ASCII, embed DejaVu Sans Mono
    // (Type0) and glyph-encode the run — base-14 Helvetica can only show ASCII, so
    // Cyrillic / accented text would otherwise render as mojibake (same defect fixed
    // for covers in 1.8.35). ASCII stays on Helvetica (`/WmF`).
    let size = spec.font_size_pt.max(4.0);
    let mut ufont = super::font::EmbeddedFont::load();
    let want_unicode = spec
        .text
        .as_ref()
        .map(|t| !t.is_ascii())
        .unwrap_or(false)
        && ufont.is_some();
    // Encode + measure BEFORE finalize (encode records glyphs, finalize consumes self).
    let uni: Option<(String, f32)> = if want_unicode {
        let text = spec.text.as_ref().unwrap();
        let f = ufont.as_mut().unwrap();
        Some((f.encode(text), f.width(text, size)))
    } else {
        None
    };

    // Image pixel size (to preserve aspect) + the shared XObject, added once.
    let image_obj = match &spec.image {
        Some(path) => {
            let (stream, w, h) = super::cover::image_xobject(path)?;
            let id = doc.document_mut().add_object(stream);
            Some((id, (w.max(1) as f32), (h.max(1) as f32)))
        }
        None => None,
    };

    // Shared ExtGState for translucency.
    let alpha = spec.opacity.clamp(0.0, 1.0);
    let mut gs = Dictionary::new();
    gs.set("Type", Object::Name(b"ExtGState".to_vec()));
    gs.set("ca", Object::Real(alpha));
    gs.set("CA", Object::Real(alpha));
    let gs_id = doc.document_mut().add_object(Object::Dictionary(gs));

    let inner = doc.document_mut();
    // Embed the Unicode font once; each stamped page references the same object.
    let uni_font_id = if want_unicode {
        Some(ufont.take().unwrap().finalize(inner))
    } else {
        None
    };
    let mut stamped = 0usize;
    for &page_no in &selected {
        let idx = (page_no - 1) as usize;
        let Some(&pid) = page_ids.get(idx) else { continue };
        // Page geometry from MediaBox.
        let (pw, ph) = page_box(inner, pid).unwrap_or((mm_to_pt(210.0), mm_to_pt(297.0)));
        let aspect = image_obj.map(|(_, w, h)| h / w);
        let ops = build_stamp_ops(spec, pw, ph, aspect, uni.as_ref().map(|(h, w)| (h.as_str(), *w)));

        // Wire resources: ExtGState/WmGS, Font/WmF (+ /WmU), XObject/WmImg.
        inner.add_graphics_state(pid, "WmGS", gs_id).map_err(Error::Lopdf)?;
        ensure_font(inner, pid, uni_font_id)?;
        if let Some((img_id, _, _)) = image_obj {
            inner.add_xobject(pid, "WmImg", img_id).map_err(Error::Lopdf)?;
        }

        // Append the stamp after existing content (collapsed to one stream).
        let mut content = inner.get_and_decode_page_content(pid).map_err(Error::Lopdf)?;
        let extra = lopdf::content::Content::decode(ops.as_bytes())
            .map_err(Error::Lopdf)?;
        content.operations.extend(extra.operations);
        let encoded = content.encode().map_err(Error::Lopdf)?;
        inner.change_page_content(pid, encoded).map_err(Error::Lopdf)?;
        stamped += 1;
    }
    Ok(stamped)
}

/// Build the `q … Q` stamp content for one page of size `pw × ph`.
/// `image_aspect` is the stamp image's height/width (None when no image).
fn build_stamp_ops(
    spec: &WatermarkSpec,
    pw: f32,
    ph: f32,
    image_aspect: Option<f32>,
    uni: Option<(&str, f32)>,
) -> String {
    let (ax, ay) = anchor(spec.position, pw, ph);
    let mut s = String::from("q\n/WmGS gs\n");

    if let Some(aspect) = image_aspect {
        // Centre the scaled image on the anchor, preserving aspect.
        let tw = pw * spec.image_scale.clamp(0.02, 1.0);
        let th = tw * aspect;
        let ix = ax - tw / 2.0;
        let iy = ay - th / 2.0;
        s.push_str(&format!("q {tw:.3} 0 0 {th:.3} {ix:.3} {iy:.3} cm /WmImg Do Q\n"));
    }

    if let Some(text) = &spec.text {
        let size = spec.font_size_pt.max(4.0);
        let (r, g, b) = spec.color;
        let rad = spec.rotation_deg.to_radians();
        let (cos, sin) = (rad.cos(), rad.sin());
        // Pick font + show-operand: embedded Unicode (hex, measured advance) or
        // base-14 Helvetica (parenthesised, rough advance) for ASCII.
        let (font_name, show, tw) = match uni {
            Some((hex, w)) => ("WmU", hex.to_string(), w),
            None => ("WmF", format!("({})", esc(text)), size * 0.5 * text.chars().count() as f32),
        };
        s.push_str(&format!("{r:.3} {g:.3} {b:.3} rg\n"));
        s.push_str(&format!(
            "BT /{font_name} {size:.1} Tf {cos:.5} {sin:.5} {nsin:.5} {cos:.5} {ax:.3} {ay:.3} Tm\n",
            nsin = -sin
        ));
        // Shift along the rotated baseline to centre the run on the anchor.
        s.push_str(&format!("{:.3} {:.3} Td {} Tj ET\n", -tw / 2.0, -size * 0.35, show));
    }

    s.push_str("Q\n");
    s
}

fn anchor(pos: WmPosition, pw: f32, ph: f32) -> (f32, f32) {
    match pos {
        WmPosition::Center => (pw / 2.0, ph / 2.0),
        WmPosition::TopLeft => (MARGIN, ph - MARGIN),
        WmPosition::TopRight => (pw - MARGIN, ph - MARGIN),
        WmPosition::BottomLeft => (MARGIN, MARGIN),
        WmPosition::BottomRight => (pw - MARGIN, MARGIN),
    }
}

fn page_box(doc: &lopdf::Document, page_id: lopdf::ObjectId) -> Option<(f32, f32)> {
    let dict = doc.get_dictionary(page_id).ok()?;
    let mb = dict.get(b"MediaBox").ok()?.as_array().ok()?;
    if mb.len() != 4 {
        return None;
    }
    let v: Vec<f32> = mb.iter().map(|o| o.as_float().unwrap_or(0.0)).collect();
    Some(((v[2] - v[0]).abs(), (v[3] - v[1]).abs()))
}

/// Ensure the page's Resources carry an inline Helvetica as `/WmF`, and — when the
/// watermark text is non-ASCII — the embedded Unicode font (`uni_font_id`) as `/WmU`.
fn ensure_font(
    doc: &mut lopdf::Document,
    page_id: lopdf::ObjectId,
    uni_font_id: Option<lopdf::ObjectId>,
) -> Result<()> {
    let res = doc
        .get_or_create_resources(page_id)
        .map_err(Error::Lopdf)?
        .as_dict_mut()
        .map_err(Error::Lopdf)?;
    if !res.has(b"Font") {
        res.set("Font", Dictionary::new());
    }
    let fonts = res.get_mut(b"Font").and_then(Object::as_dict_mut).map_err(Error::Lopdf)?;
    if !fonts.has(b"WmF") {
        let mut helv = Dictionary::new();
        helv.set("Type", "Font");
        helv.set("Subtype", "Type1");
        helv.set("BaseFont", "Helvetica");
        fonts.set("WmF", Object::Dictionary(helv));
    }
    if let Some(fid) = uni_font_id {
        if !fonts.has(b"WmU") {
            fonts.set("WmU", Object::Reference(fid));
        }
    }
    Ok(())
}

fn esc(s: &str) -> String {
    let mut out = String::with_capacity(s.len());
    for ch in s.chars() {
        if matches!(ch, '(' | ')' | '\\') {
            out.push('\\');
        }
        out.push(ch);
    }
    out
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::pdf::PdfDoc;
    use lopdf::Stream;

    /// A small helper PDF whose pages carry a trivial content stream.
    fn doc_with_content(n: usize) -> PdfDoc {
        let mut pdf = PdfDoc::load_mem(&crate::pdf::test_support::minimal_pdf(n, 300.0, 400.0)).unwrap();
        let ids = pdf.page_ids().to_vec();
        let inner = pdf.document_mut();
        for pid in ids {
            let cid = inner.add_object(Stream::new(Dictionary::new(), b"0 0 0 rg\n".to_vec()));
            if let Ok(Object::Dictionary(p)) = inner.get_object_mut(pid) {
                p.set("Contents", cid);
            }
        }
        pdf
    }

    #[test]
    fn position_parses() {
        assert_eq!(WmPosition::parse("center"), Some(WmPosition::Center));
        assert_eq!(WmPosition::parse("bottom-right"), Some(WmPosition::BottomRight));
        assert_eq!(WmPosition::parse("TL"), Some(WmPosition::TopLeft));
        assert!(WmPosition::parse("sideways").is_none());
    }

    #[test]
    fn anchors_land_in_the_right_corners() {
        let (pw, ph) = (300.0, 400.0);
        assert_eq!(anchor(WmPosition::Center, pw, ph), (150.0, 200.0));
        assert_eq!(anchor(WmPosition::TopRight, pw, ph), (pw - MARGIN, ph - MARGIN));
        assert_eq!(anchor(WmPosition::BottomLeft, pw, ph), (MARGIN, MARGIN));
    }

    #[test]
    fn empty_spec_errors() {
        let mut pdf = doc_with_content(1);
        assert!(apply_watermark(&mut pdf, &WatermarkSpec::default()).is_err());
    }

    #[test]
    fn stamps_text_on_all_pages_and_round_trips() {
        let mut pdf = doc_with_content(3);
        let spec = WatermarkSpec {
            text: Some("DRAFT".into()),
            ..Default::default()
        };
        let n = apply_watermark(&mut pdf, &spec).unwrap();
        assert_eq!(n, 3);
        // every page now contains the DRAFT Tj + the gs reference
        let reloaded = PdfDoc::load_mem(&pdf.to_bytes().unwrap()).unwrap();
        assert_eq!(reloaded.page_count(), 3);
        for pid in reloaded.page_ids() {
            let c = reloaded.document().get_and_decode_page_content(*pid).unwrap();
            assert!(
                c.operations.iter().any(|o| o.operator == "gs"),
                "page carries the watermark ExtGState"
            );
            assert!(
                c.operations.iter().any(|o| o.operator == "Tj"),
                "page carries the watermark text"
            );
        }
    }

    #[test]
    fn page_range_limits_the_stamp() {
        let mut pdf = doc_with_content(4);
        let spec = WatermarkSpec {
            text: Some("X".into()),
            pages: Some(PageSpec::parse("2-3").unwrap()),
            ..Default::default()
        };
        assert_eq!(apply_watermark(&mut pdf, &spec).unwrap(), 2);
        let count_gs = |pdf: &PdfDoc, idx: usize| {
            let pid = pdf.page_ids()[idx];
            pdf.document()
                .get_and_decode_page_content(pid)
                .unwrap()
                .operations
                .iter()
                .filter(|o| o.operator == "gs")
                .count()
        };
        assert_eq!(count_gs(&pdf, 0), 0, "page 1 untouched");
        assert_eq!(count_gs(&pdf, 1), 1, "page 2 stamped");
        assert_eq!(count_gs(&pdf, 3), 0, "page 4 untouched");
    }

    #[test]
    fn ascii_uses_helvetica_unicode_uses_embedded_font() {
        let spec = WatermarkSpec {
            text: Some("DRAFT".into()),
            ..Default::default()
        };
        // ASCII → base-14 Helvetica, literal parenthesised string.
        let ascii = build_stamp_ops(&spec, 300.0, 400.0, None, None);
        assert!(ascii.contains("/WmF "), "ASCII stamp uses /WmF: {ascii}");
        assert!(ascii.contains("(DRAFT)"), "ASCII stamp shows literal text");
        assert!(!ascii.contains("/WmU"));

        // Non-ASCII → embedded Unicode font, glyph-hex show string.
        let uni = build_stamp_ops(&spec, 300.0, 400.0, None, Some(("<04220430>", 12.0)));
        assert!(uni.contains("/WmU "), "unicode stamp uses /WmU: {uni}");
        assert!(uni.contains("<04220430> Tj"), "unicode stamp shows glyph hex");
        assert!(!uni.contains("/WmF"));
    }

    #[test]
    fn cyrillic_watermark_registers_unicode_font_resource() {
        // Only meaningful when the bundled DejaVu font is available; otherwise the
        // watermark silently falls back to Helvetica and this assertion is skipped.
        if super::super::font::EmbeddedFont::load().is_none() {
            return;
        }
        let mut pdf = doc_with_content(1);
        let spec = WatermarkSpec {
            text: Some("ЧЕРНОВИК".into()), // "DRAFT" in Russian
            ..Default::default()
        };
        assert_eq!(apply_watermark(&mut pdf, &spec).unwrap(), 1);
        let reloaded = PdfDoc::load_mem(&pdf.to_bytes().unwrap()).unwrap();
        let pid = reloaded.page_ids()[0];
        let content = reloaded.document().get_and_decode_page_content(pid).unwrap();
        // The Unicode path shows a glyph-hex string; the Helvetica fallback would
        // show a literal string. Assert a Tj carries a Hexadecimal operand.
        let hex_tj = content.operations.iter().any(|op| {
            op.operator == "Tj"
                && matches!(
                    op.operands.first(),
                    Some(Object::String(_, lopdf::StringFormat::Hexadecimal))
                )
        });
        assert!(hex_tj, "Cyrillic watermark is drawn with the embedded Unicode font (hex Tj)");
    }
}