use nom::{
    bytes::complete::{tag, take, take_until},
    character::complete,
    multi::separated_list1,
    IResult,
};

#[derive(Debug)]
struct Record {
    springs: String,
    groups: Vec<u32>,
}

fn parse_line(input: &str) -> IResult<&str, Record> {
    let (input, springs) = take_until(" ")(input)?;
    let (input, _) = take(1 as usize)(input)?;
    let (input, groups) = separated_list1(tag(","), complete::u32)(input)?;
    Ok((input, Record { springs: springs.to_string(), groups }))
}

fn total_arrangements(springs: &str, groups: &[u32], cache: &mut Vec<Vec<Option<u64>>>) -> u64 {
    // If there isn't enough space for remaining groups return 0
    // Else if the counts are empty return 1 (You found one!)
    if springs.len() < groups.iter().sum::<u32>() as usize + groups.len() {
        return 0;
    } else if groups.is_empty() {
        return if springs.contains("#") { 0 } else { 1 }
    };

    // If in cache return
    if let Some(cached) = cache[groups.len()-1][springs.len()-1] {
        return cached;
    };

    // Total arrangements = 0
    let mut total = 0;

    // Assume first value is a "." whether it's a "." or a "?" and don't reduce the count
    if springs.chars().nth(0).unwrap() != '#' {
        total += total_arrangements(&springs[1..], groups, cache);
    }

    // Determine arrangements if first groups[0] could all be # then there's a . (remove groups[0])
    let group_size = groups[0] as usize;
    if !&springs[..group_size].contains(".") && springs.chars().nth(group_size).unwrap() != '#' {
        total += total_arrangements(&springs[group_size+1..], &groups[1..], cache)
    }

    // Add arrangements to cache for that size of groups and springs
    cache[groups.len() - 1][springs.len() - 1] = Some(total);
    total
}

pub fn part1(input: &str) -> String {
    input
        .lines()
        .filter_map(|line| parse_line(line).ok().and_then(|x| Some(x.1)))
        .map(|Record {springs, groups}| {
            let mut cache = vec![vec![None; springs.len() + 1]; groups.len()];
            total_arrangements(&(springs + "."), &groups, &mut cache)
        })
        .sum::<u64>()
        .to_string()
}

pub fn part2(input: &str) -> String {
    input
        .lines()
        .filter_map(|line| parse_line(line).ok().and_then(|x| Some(x.1)))
        .map(|Record {springs, groups}| {
            let mut new_springs = (springs + &"?").repeat(5);
            new_springs.pop();
            let new_groups = std::iter::repeat(groups).take(5).flatten().collect();
            Record {springs: new_springs, groups: new_groups}
        })
        .enumerate()
        .map(|(_, Record {springs, groups})| {
            let mut cache = vec![vec![None; springs.len() + 1]; groups.len()];
            total_arrangements(&(springs + "."), &groups, &mut cache)
        })
        .sum::<u64>()
        .to_string()
}

pub mod prelude {
    pub use super::part1;
    pub use super::part2;
}
