kindexr/src/enrich/parser.rs

use std::sync::OnceLock;

use regex::Regex;

pub struct ParsedTitle {
    pub clean: String,
    pub season: Option<i32>,
    pub episode: Option<i32>,
    pub year: Option<i32>,
    pub quality: Option<String>,
    pub source: Option<String>,
}

static SE_RE: OnceLock<Regex> = OnceLock::new();
static ALT_SE_RE: OnceLock<Regex> = OnceLock::new();
static YEAR_RE: OnceLock<Regex> = OnceLock::new();
static QUALITY_RE: OnceLock<Regex> = OnceLock::new();
static SOURCE_RE: OnceLock<Regex> = OnceLock::new();

fn se_re() -> &'static Regex {
    SE_RE.get_or_init(|| Regex::new(r"(?i)\bS(\d{1,2})E(\d{1,4})\b").unwrap())
}

fn alt_se_re() -> &'static Regex {
    ALT_SE_RE.get_or_init(|| Regex::new(r"\b(\d{1,2})x(\d{2,4})\b").unwrap())
}

fn year_re() -> &'static Regex {
    YEAR_RE.get_or_init(|| Regex::new(r"\b((?:19|20)\d{2})\b").unwrap())
}

fn quality_re() -> &'static Regex {
    QUALITY_RE.get_or_init(|| Regex::new(r"(?i)\b(2160p|4K|UHD|1080p|720p|576p|480p)\b").unwrap())
}

fn source_re() -> &'static Regex {
    SOURCE_RE.get_or_init(|| {
        Regex::new(r"(?i)\b(WEB-DL|WEBRip|BluRay|BDRip|HDTV|DVDRip|REMUX)\b").unwrap()
    })
}

pub fn parse(raw: &str) -> ParsedTitle {
    let s: String = raw.chars().map(|c| if c == '.' || c == '_' { ' ' } else { c }).collect();

    let (season, episode) = extract_season_ep(&s);
    let year = extract_year(&s);
    let quality = quality_re().captures(&s).and_then(|c| c.get(1)).map(|m| m.as_str().to_owned());
    let source = source_re().captures(&s).and_then(|c| c.get(1)).map(|m| m.as_str().to_owned());

    let cut = find_cut(&s, season, quality.is_some(), source.is_some());
    let clean = s[..cut].split_whitespace().collect::<Vec<_>>().join(" ");

    ParsedTitle { clean, season, episode, year, quality, source }
}

fn extract_season_ep(s: &str) -> (Option<i32>, Option<i32>) {
    if let Some(caps) = se_re().captures(s) {
        return (
            caps.get(1).and_then(|m| m.as_str().parse().ok()),
            caps.get(2).and_then(|m| m.as_str().parse().ok()),
        );
    }
    if let Some(caps) = alt_se_re().captures(s) {
        return (
            caps.get(1).and_then(|m| m.as_str().parse().ok()),
            caps.get(2).and_then(|m| m.as_str().parse().ok()),
        );
    }
    (None, None)
}

fn extract_year(s: &str) -> Option<i32> {
    year_re()
        .captures_iter(s)
        .filter_map(|c| c.get(1)?.as_str().parse::<i32>().ok())
        .find(|&y| y >= 1900 && y <= 2099)
}

fn find_cut(s: &str, season: Option<i32>, has_quality: bool, has_source: bool) -> usize {
    let mut cut = s.len();
    if season.is_some() {
        if let Some(m) = se_re().find(s) { cut = cut.min(m.start()); }
        if let Some(m) = alt_se_re().find(s) { cut = cut.min(m.start()); }
    }
    if has_quality {
        if let Some(m) = quality_re().find(s) { cut = cut.min(m.start()); }
    }
    if has_source {
        if let Some(m) = source_re().find(s) { cut = cut.min(m.start()); }
    }
    cut
}

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

    #[test]
    fn season_ep() {
        let p = parse("Breaking.Bad.S01E07.720p.BluRay.x264");
        assert_eq!(p.season, Some(1));
        assert_eq!(p.episode, Some(7));
        assert_eq!(p.clean, "Breaking Bad");
        assert_eq!(p.quality.as_deref(), Some("720p"));
        assert_eq!(p.source.as_deref(), Some("BluRay"));
    }

    #[test]
    fn alt_season_ep() {
        let p = parse("Show.Name.1x03.HDTV");
        assert_eq!(p.season, Some(1));
        assert_eq!(p.episode, Some(3));
        assert_eq!(p.source.as_deref(), Some("HDTV"));
    }

    #[test]
    fn movie_with_year() {
        let p = parse("The.Dark.Knight.2008.1080p.WEB-DL");
        assert_eq!(p.season, None);
        assert_eq!(p.year, Some(2008));
        assert_eq!(p.quality.as_deref(), Some("1080p"));
        // Clean title includes year since year alone doesn't cut
        assert!(p.clean.contains("Dark Knight"));
    }

    #[test]
    fn no_markers() {
        let p = parse("Just a plain title");
        assert_eq!(p.clean, "Just a plain title");
        assert_eq!(p.season, None);
        assert_eq!(p.quality, None);
    }
}