//! Secret masking for job logs. Replaces every secret value, and its standard base64 encoding,
//! with `***`. [`MaskStream`] does the same over chunks, so a secret split across two chunks is
//! still masked.

use std::collections::BTreeSet;
use std::fmt;

use base64::Engine as _;
use base64::engine::general_purpose::STANDARD;

/// The text that replaces every match.
pub const MASK: &str = "***";
/// Values shorter than this (in chars) are not masked, since they would hide ordinary output.
pub const MIN_SECRET_CHARS: usize = 4;

/// Masks a fixed set of secret values.
#[derive(Clone)]
pub struct Masker {
    /// Byte patterns (raw values and their base64), longest first, deduplicated.
    patterns: Vec<Vec<u8>>,
}

impl fmt::Debug for Masker {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("Masker")
            .field("patterns", &self.patterns.len())
            .finish()
    }
}

impl Masker {
    /// Builds a masker. Values shorter than [`MIN_SECRET_CHARS`] are ignored.
    pub fn new(secrets: impl IntoIterator<Item = String>) -> Self {
        let mut set: BTreeSet<Vec<u8>> = BTreeSet::new();
        for value in secrets {
            if value.chars().count() < MIN_SECRET_CHARS {
                continue;
            }
            let encoded = STANDARD.encode(value.as_bytes());
            set.insert(value.into_bytes());
            set.insert(encoded.into_bytes());
        }
        let mut patterns: Vec<Vec<u8>> = set.into_iter().collect();
        patterns.sort_by_key(|p| std::cmp::Reverse(p.len()));
        Self { patterns }
    }

    /// Masks every occurrence in `text`.
    pub fn mask(&self, text: &str) -> String {
        let (out, _) = scan(&self.patterns, text.as_bytes(), true);
        bytes_to_string(out)
    }

    /// Starts a streaming masker that owns a copy of this masker.
    pub fn stream(&self) -> MaskStream {
        MaskStream::new(self.clone())
    }
}

/// Masks a byte stream that arrives in chunks.
#[derive(Debug, Clone)]
pub struct MaskStream {
    masker: Masker,
    /// Bytes held back because they may be the start of a secret. Never more than the longest
    /// pattern minus one.
    held: Vec<u8>,
}

impl MaskStream {
    pub fn new(masker: Masker) -> Self {
        Self {
            masker,
            held: Vec::new(),
        }
    }

    /// Feeds one chunk. Returns the masked bytes that are now safe to emit.
    pub fn push(&mut self, chunk: &[u8]) -> Vec<u8> {
        self.held.extend_from_slice(chunk);
        let (out, rest) = scan(&self.masker.patterns, &self.held, false);
        self.held = rest;
        out
    }

    /// Ends the stream and returns whatever is still held back, masked.
    pub fn finish(&mut self) -> Vec<u8> {
        let held = std::mem::take(&mut self.held);
        let (out, _) = scan(&self.masker.patterns, &held, true);
        out
    }
}

/// Walks `input` and writes masked bytes. A match is replaced by [`MASK`]. When `last` is false,
/// the scan stops at the first position where a pattern could still complete in later input, and
/// returns that unprocessed tail. When `last` is true, the whole input is processed.
fn scan(patterns: &[Vec<u8>], input: &[u8], last: bool) -> (Vec<u8>, Vec<u8>) {
    let mut out = Vec::with_capacity(input.len());
    let mut i = 0;
    while i < input.len() {
        let tail = &input[i..];
        if let Some(p) = patterns.iter().find(|p| tail.starts_with(p)) {
            out.extend_from_slice(MASK.as_bytes());
            i += p.len();
            continue;
        }
        if !last
            && patterns
                .iter()
                .any(|p| p.len() > tail.len() && p.starts_with(tail))
        {
            return (out, tail.to_vec());
        }
        out.push(input[i]);
        i += 1;
    }
    (out, Vec::new())
}

fn bytes_to_string(bytes: Vec<u8>) -> String {
    match String::from_utf8(bytes) {
        Ok(s) => s,
        Err(e) => String::from_utf8_lossy(e.as_bytes()).into_owned(),
    }
}

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

    fn masker(values: &[&str]) -> Masker {
        Masker::new(values.iter().map(|v| v.to_string()))
    }

    fn stream_all(m: &Masker, chunks: &[&[u8]]) -> String {
        let mut s = m.stream();
        let mut out = Vec::new();
        for c in chunks {
            out.extend(s.push(c));
        }
        out.extend(s.finish());
        String::from_utf8(out).unwrap()
    }

    #[test]
    fn masks_plain_values() {
        let m = masker(&["hunter2-secret"]);
        assert_eq!(m.mask("token=hunter2-secret ok"), "token=*** ok");
        assert_eq!(
            m.mask("two hunter2-secret and hunter2-secret"),
            "two *** and ***"
        );
        assert_eq!(m.mask("nothing here"), "nothing here");
    }

    #[test]
    fn masks_base64_encoding() {
        let secret = "hunter2-secret";
        let encoded = STANDARD.encode(secret.as_bytes());
        let m = masker(&[secret]);
        assert_eq!(m.mask(&format!("leak: {encoded}")), "leak: ***");
        let other = masker(&["abcdefgh"]);
        let enc = STANDARD.encode("abcdefgh");
        assert_eq!(other.mask(&enc), "***");
    }

    #[test]
    fn short_values_are_ignored() {
        let m = masker(&["abc", "", "ab"]);
        assert_eq!(m.mask("abc ab"), "abc ab");
        // 4 chars is the minimum that is masked.
        let m = masker(&["abcd"]);
        assert_eq!(m.mask("xabcdx"), "x***x");
    }

    #[test]
    fn longest_match_wins() {
        let m = masker(&["secret-value", "secret"]);
        assert_eq!(m.mask("a secret-value b"), "a *** b");
    }

    #[test]
    fn empty_secret_set_is_a_no_op() {
        let m = Masker::new(std::iter::empty());
        assert_eq!(m.mask("anything"), "anything");
        assert_eq!(stream_all(&m, &[b"anything"]), "anything");
    }

    #[test]
    fn secret_split_across_two_chunks() {
        let m = masker(&["hunter2-secret"]);
        let out = stream_all(&m, &[b"pre hunter2-", b"secret post"]);
        assert_eq!(out, "pre *** post");
    }

    #[test]
    fn secret_split_one_byte_at_a_time() {
        let m = masker(&["hunter2-secret"]);
        let text = b"a hunter2-secret b hunter2-secret c";
        let chunks: Vec<&[u8]> = text.chunks(1).collect();
        assert_eq!(stream_all(&m, &chunks), "a *** b *** c");
    }

    #[test]
    fn base64_split_across_chunks() {
        let secret = "hunter2-secret";
        let encoded = STANDARD.encode(secret);
        let m = masker(&[secret]);
        let (a, b) = encoded.split_at(5);
        let out = stream_all(
            &m,
            &[format!("x{a}").as_bytes(), format!("{b}y").as_bytes()],
        );
        assert_eq!(out, "x***y");
    }

    #[test]
    fn held_bytes_are_bounded() {
        let m = masker(&["0123456789"]);
        let mut s = m.stream();
        let out = s.push(b"0123456");
        assert!(out.is_empty());
        assert!(s.held.len() <= 9);
        let out = s.push(b"x");
        assert_eq!(out, b"0123456x".to_vec());
    }

    #[test]
    fn partial_prefix_that_does_not_complete_is_emitted() {
        let m = masker(&["hunter2-secret"]);
        let out = stream_all(&m, &[b"hunt", b"er is here"]);
        assert_eq!(out, "hunter is here");
    }

    #[test]
    fn finish_flushes_held_partial_match() {
        let m = masker(&["hunter2-secret"]);
        let mut s = m.stream();
        let mut out = s.push(b"end hunter2-se");
        out.extend(s.finish());
        assert_eq!(String::from_utf8(out).unwrap(), "end hunter2-se");
    }

    #[test]
    fn multibyte_text_survives() {
        let m = masker(&["hunter2-secret"]);
        assert_eq!(m.mask("héllo 🙂 hunter2-secret ✓"), "héllo 🙂 *** ✓");
        let out = stream_all(&m, &["héllo 🙂 ".as_bytes(), "hunter2-secret ✓".as_bytes()]);
        assert_eq!(out, "héllo 🙂 *** ✓");
    }

    #[test]
    fn debug_does_not_print_values() {
        let m = masker(&["hunter2-secret"]);
        assert!(!format!("{m:?}").contains("hunter2"));
    }
}