use std::collections::HashMap;

use itertools::Itertools;

pub fn process_part1(input: &str) -> u64 {
    let (patterns, designs) = input.split_once("\n\n").unwrap();
    let patterns = patterns.split(", ").collect_vec();
    designs
        .lines()
        .map(|design| {
            if possible_design(design, &patterns) {
                1
            } else {
                0
            }
        })
        .sum()
}

fn possible_design(design: &str, patterns: &[&str]) -> bool {
    for pattern in patterns {
        if design.starts_with(pattern) && design.len() == pattern.len() {
            return true;
        } else if design.starts_with(pattern) {
            if let Some(stripped) = design.strip_prefix(pattern) {
                if possible_design(stripped, patterns) {
                    return true;
                }
            }
        }
    }
    false
}

fn possible_designs(design: &str, patterns: &[&str], memo: &mut HashMap<String, u64>) -> u64 {
    if design.is_empty() {
        return 0;
    }
    let mut num_possible_designs = 0;
    if let Some(count) = memo.get(design) {
        return *count;
    }
    for pattern in patterns {
        if design == *pattern {
            memo.insert(design.to_string(), 1);
            num_possible_designs += 1;
        }
        if design.starts_with(pattern) {
            if let Some(stripped) = design.strip_prefix(pattern) {
                let for_stripped = possible_designs(stripped, patterns, memo);
                memo.entry(stripped.to_string()).or_insert(for_stripped);
                num_possible_designs += for_stripped;
            }
        }
    }
    memo.insert(design.to_string(), num_possible_designs);
    num_possible_designs
}

pub fn process_part2(input: &str) -> u64 {
    let (patterns, designs) = input.split_once("\n\n").unwrap();
    let patterns = patterns.split(", ").collect_vec();
    let mut memo = HashMap::new();
    designs
        .lines()
        .map(|design| possible_designs(design, &patterns, &mut memo))
        .sum()
}

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

    const INPUT: &str = "r, wr, b, g, bwu, rb, gb, br

brwrr
bggr
gbbr
rrbgbr
ubwu
bwurrg
brgr
bbrgwb";

    #[test]
    fn part1() {
        let result = process_part1(INPUT);
        assert_eq!(result, 6);
    }

    #[test]
    fn part2() {
        let result = process_part2(INPUT);
        assert_eq!(result, 16);
    }
}