rsomics-seq-grep 0.1.1

use std::collections::HashSet;
use std::fs;
use std::path::Path;

use regex::bytes::Regex;
use rsomics_common::{Result, RsomicsError};

/// One compiled matcher per `-p`/`-f` pattern set. The three modes mirror
/// `seqkit grep`'s three matching strategies, which are mutually exclusive
/// per invocation (chosen once from `-r`/`-d`/`-s`, not per pattern).
pub enum PatternSet {
    /// `-p`/`-f` without `-r`/`-d`, against ID or full name: whole-string equality.
    Exact(HashSet<Vec<u8>>),
    /// `-s` without `-r`/`-d`: substring search on the sequence.
    Substring(Vec<Vec<u8>>),
    /// `-r` (any target) or `-d` (sequence only, IUPAC-expanded first): unanchored regex.
    Regex(Vec<Regex>),
}

impl PatternSet {
    pub fn is_match(&self, target: &[u8], ignore_case: bool) -> bool {
        match self {
            PatternSet::Exact(set) => {
                if ignore_case {
                    set.contains(&target.to_ascii_lowercase())
                } else {
                    set.contains(target)
                }
            }
            PatternSet::Substring(pats) => {
                let lowered;
                let t: &[u8] = if ignore_case {
                    lowered = target.to_ascii_lowercase();
                    &lowered
                } else {
                    target
                };
                pats.iter().any(|p| {
                    if p.is_empty() {
                        t.is_empty()
                    } else {
                        memchr::memmem::find(t, p).is_some()
                    }
                })
            }
            PatternSet::Regex(res) => res.iter().any(|re| re.is_match(target)),
        }
    }
}

pub struct BuildOpts {
    pub use_regexp: bool,
    pub degenerate: bool,
    pub by_seq: bool,
    pub ignore_case: bool,
}

/// Build the matcher for one grep invocation. Precedence matches upstream:
/// `-d` and `-r` both compile to regex (an empty pattern becomes `^$`, an
/// exact-empty-string match, rather than the "matches everywhere" empty
/// regex); otherwise sequence search is substring, ID/name search is exact.
pub fn build(patterns: &[String], opts: &BuildOpts) -> Result<PatternSet> {
    if opts.use_regexp && opts.degenerate {
        return Err(RsomicsError::InvalidInput(
            "cannot give both -d/--degenerate and -r/--use-regexp".into(),
        ));
    }

    if opts.use_regexp || opts.degenerate {
        let mut res = Vec::with_capacity(patterns.len());
        for p in patterns {
            let expanded = if opts.degenerate {
                degenerate_to_regex(p)
            } else if p.is_empty() {
                "^$".to_string()
            } else {
                p.clone()
            };
            let source = if opts.ignore_case {
                format!("(?i){expanded}")
            } else {
                expanded
            };
            let re = Regex::new(&source)
                .map_err(|e| RsomicsError::InvalidInput(format!("invalid regex {p:?}: {e}")))?;
            res.push(re);
        }
        return Ok(PatternSet::Regex(res));
    }

    if opts.by_seq {
        for p in patterns {
            if !is_legal_seq_pattern(p.as_bytes()) {
                return Err(RsomicsError::InvalidInput(format!(
                    "illegal DNA/RNA/Protein sequence: {p}"
                )));
            }
        }
        let subs = patterns
            .iter()
            .map(|p| {
                let bytes = p.as_bytes().to_vec();
                if opts.ignore_case {
                    bytes.to_ascii_lowercase()
                } else {
                    bytes
                }
            })
            .collect();
        return Ok(PatternSet::Substring(subs));
    }

    let set = patterns
        .iter()
        .map(|p| {
            let bytes = p.as_bytes().to_vec();
            if opts.ignore_case {
                bytes.to_ascii_lowercase()
            } else {
                bytes
            }
        })
        .collect();
    Ok(PatternSet::Exact(set))
}

const fn alphabet_set(bytes: &[u8]) -> [bool; 256] {
    let mut set = [false; 256];
    let mut i = 0;
    while i < bytes.len() {
        set[bytes[i] as usize] = true;
        i += 1;
    }
    set
}

// The valid-byte set of each seqkit alphabet (bio/seq/alphabet.go): letters +
// gap (` -.`) + ambiguous codes, all case variants listed explicitly.
const DNA_REDUNDANT: [bool; 256] = alphabet_set(b"acgtryswkmbdhvACGTRYSWKMBDHV -.nN.");
const RNA_REDUNDANT: [bool; 256] = alphabet_set(b"acguryswkmbdhvACGURYSWKMBDHV -.nN");
const PROTEIN: [bool; 256] =
    alphabet_set(b"abcdefghijklmnopqrstuvwyzABCDEFGHIJKLMNOPQRSTUVWYZ -xX*_.");

/// A `-s`/`-f` literal pattern is legal only if every byte is valid in one of
/// seqkit's DNAredundant / RNAredundant / Protein alphabets — mirroring
/// `grep.go`'s `IsValid` triple-check, which rejects e.g. a digit in `-s -p T1A`.
/// An empty pattern is vacuously legal, matching seqkit.
fn is_legal_seq_pattern(p: &[u8]) -> bool {
    p.iter().all(|&b| DNA_REDUNDANT[b as usize])
        || p.iter().all(|&b| RNA_REDUNDANT[b as usize])
        || p.iter().all(|&b| PROTEIN[b as usize])
}

/// `seqkit`'s `Degenerate2Regexp` over `DegenerateBaseMapNucl` (`bio/seq/seq.go`):
/// each IUPAC nucleotide letter expands to the regex class of the bases it stands
/// for; every other byte — a plain A/C/G/T/U or a non-IUPAC char — passes through
/// literally, so `-d -p 'RY9SW'` runs (the `9` becomes a literal) rather than
/// erroring. Protein degenerate codes are not covered.
fn degenerate_to_regex(p: &str) -> String {
    if p.is_empty() {
        return "^$".to_string();
    }
    let mut out = String::with_capacity(p.len() * 2);
    for b in p.bytes() {
        match b {
            b'A' => out.push('A'),
            b'T' | b'U' => out.push_str("[TU]"),
            b'C' => out.push('C'),
            b'G' => out.push('G'),
            b'R' => out.push_str("[AG]"),
            b'Y' => out.push_str("[CTU]"),
            b'M' => out.push_str("[AC]"),
            b'K' => out.push_str("[GTU]"),
            b'S' => out.push_str("[CG]"),
            b'W' => out.push_str("[ATU]"),
            b'H' => out.push_str("[ACTU]"),
            b'B' => out.push_str("[CGTU]"),
            b'V' => out.push_str("[ACG]"),
            b'D' => out.push_str("[AGTU]"),
            b'N' => out.push_str("[ACGTU]"),
            b'a' => out.push('a'),
            b't' | b'u' => out.push_str("[tu]"),
            b'c' => out.push('c'),
            b'g' => out.push('g'),
            b'r' => out.push_str("[ag]"),
            b'y' => out.push_str("[ctu]"),
            b'm' => out.push_str("[ac]"),
            b'k' => out.push_str("[gtu]"),
            b's' => out.push_str("[cg]"),
            b'w' => out.push_str("[atu]"),
            b'h' => out.push_str("[actu]"),
            b'b' => out.push_str("[cgtu]"),
            b'v' => out.push_str("[acg]"),
            b'd' => out.push_str("[agtu]"),
            b'n' => out.push_str("[acgtu]"),
            other => out.push(other as char),
        }
    }
    out
}

/// Load patterns from a plain-text file, one per line. Trailing `\r\n`/`\n`
/// is stripped; blank lines are skipped (matches `seqkit`'s `breader`
/// default line reader plus its own empty-line skip in `grep.go`).
pub fn load_pattern_file(path: &Path) -> Result<Vec<String>> {
    let raw = fs::read(path)
        .map_err(|e| RsomicsError::InvalidInput(format!("{}: {e}", path.display())))?;
    let text = String::from_utf8(raw)
        .map_err(|e| RsomicsError::InvalidInput(format!("{}: {e}", path.display())))?;
    Ok(text
        .split('\n')
        .map(|line| line.strip_suffix('\r').unwrap_or(line))
        .filter(|line| !line.is_empty())
        .map(str::to_owned)
        .collect())
}

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

    fn opts(use_regexp: bool, degenerate: bool, by_seq: bool, ignore_case: bool) -> BuildOpts {
        BuildOpts {
            use_regexp,
            degenerate,
            by_seq,
            ignore_case,
        }
    }

    #[test]
    fn exact_match_is_whole_string() {
        let set = build(&["seq1".to_string()], &opts(false, false, false, false)).unwrap();
        assert!(set.is_match(b"seq1", false));
        assert!(!set.is_match(b"seq1_dup", false));
    }

    #[test]
    fn substring_matches_partial() {
        let set = build(&["TTTT".to_string()], &opts(false, false, true, false)).unwrap();
        assert!(set.is_match(b"AATTTTAA", false));
        assert!(!set.is_match(b"AACCCCAA", false));
    }

    #[test]
    fn empty_substring_pattern_only_matches_empty_target() {
        let set = build(&[String::new()], &opts(false, false, true, false)).unwrap();
        assert!(set.is_match(b"", false));
        assert!(!set.is_match(b"ACGT", false));
    }

    #[test]
    fn degenerate_expands_iupac() {
        let set = build(&["RYSW".to_string()], &opts(false, true, true, false)).unwrap();
        assert!(set.is_match(b"ACGT", false));
        assert!(!set.is_match(b"CCCC", false));
    }

    #[test]
    fn degenerate_passes_non_iupac_through() {
        // A non-IUPAC byte becomes a literal regex char (no error).
        let set = build(&["RY9SW".to_string()], &opts(false, true, true, false)).unwrap();
        assert!(set.is_match(b"AC9CA", false));
        assert!(!set.is_match(b"ACGCA", false));
    }

    #[test]
    fn by_seq_rejects_illegal_alphabet() {
        assert!(build(&["T1A".to_string()], &opts(false, false, true, false)).is_err());
    }

    #[test]
    fn by_seq_accepts_iupac_and_protein() {
        for p in ["ATGC", "RYSWKMBDHVN", "MEK", "AT.GC", ""] {
            assert!(
                build(&[p.to_string()], &opts(false, false, true, false)).is_ok(),
                "should accept {p:?}"
            );
        }
    }

    #[test]
    fn regex_is_unanchored() {
        let set = build(&["seq1".to_string()], &opts(true, false, false, false)).unwrap();
        assert!(set.is_match(b"seq1_dup", false));
    }

    #[test]
    fn ignore_case_lowercases_both_sides() {
        let set = build(&["SEQ1".to_string()], &opts(false, false, false, true)).unwrap();
        assert!(set.is_match(b"seq1", true));
    }

    #[test]
    fn empty_regex_pattern_is_anchored_empty() {
        let set = build(&[String::new()], &opts(true, false, false, false)).unwrap();
        assert!(set.is_match(b"", false));
        assert!(!set.is_match(b"anything", false));
    }
}