rsomics-seq-grep 0.1.0

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 {
        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))
}

/// `seqkit`'s `DegenerateBaseMapNucl` (`bio/seq/seq.go`): each IUPAC nucleotide
/// letter expands to the regex class of the bases it stands for; a plain
/// A/C/G/T/U passes through as itself. Protein degenerate codes are not
/// covered — `-d` is documented here as nucleotide-only.
fn degenerate_to_regex(p: &str) -> Result<String> {
    if p.is_empty() {
        return Ok("^$".to_string());
    }
    let mut out = String::with_capacity(p.len() * 2);
    for b in p.bytes() {
        let cls = match b {
            b'A' => "A",
            b'T' | b'U' => "[TU]",
            b'C' => "C",
            b'G' => "G",
            b'R' => "[AG]",
            b'Y' => "[CTU]",
            b'M' => "[AC]",
            b'K' => "[GTU]",
            b'S' => "[CG]",
            b'W' => "[ATU]",
            b'H' => "[ACTU]",
            b'B' => "[CGTU]",
            b'V' => "[ACG]",
            b'D' => "[AGTU]",
            b'N' => "[ACGTU]",
            b'a' => "a",
            b't' | b'u' => "[tu]",
            b'c' => "c",
            b'g' => "g",
            b'r' => "[ag]",
            b'y' => "[ctu]",
            b'm' => "[ac]",
            b'k' => "[gtu]",
            b's' => "[cg]",
            b'w' => "[atu]",
            b'h' => "[actu]",
            b'b' => "[cgtu]",
            b'v' => "[acg]",
            b'd' => "[agtu]",
            b'n' => "[acgtu]",
            _ => {
                return Err(RsomicsError::InvalidInput(format!(
                    "-d/--degenerate: `{p}` is not a valid IUPAC nucleotide motif (byte {b:#04x} unrecognised)"
                )));
            }
        };
        out.push_str(cls);
    }
    Ok(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 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));
    }
}