use aoc_runner_derive::{aoc, aoc_generator};
use itertools::Itertools;
use std::collections::HashMap;
use std::slice;

#[aoc_generator(day12)]
fn parse(input: &str) -> Vec<(Vec<Condition>, Vec<usize>)> {
    let input: Vec<(Vec<Condition>, Vec<usize>)> = input.split('\n')
        .map(|line| {
            line.split(' ')
                .collect_tuple()
                .unwrap()
        }).map(|(springs, groups)| {
            let springs: Vec<Condition> = springs.chars()
                .map(|char| {
                    char.to_condition()
                }).collect();

            let groups: Vec<_> = groups.split(',')
                .map(|num| num.parse::<usize>().expect("Failed to parse group len"))
                .collect();
            (springs, groups)
        }).collect();
    input
}

fn possible_count_cache(springs: &[Condition], groups: &[usize], cache: &mut HashMap<(usize, usize), usize>) -> usize {
    if let Some(count) = cache.get(&(springs.len(), groups.len())) {
        return *count;
    }

    let mut count = 0;

    if groups.is_empty() {
        count = if springs.contains(&Condition::Bad) { 0 } else { 1 };
        cache.insert((springs.len(), groups.len()), count);

        return count;
    }

    for i in 0..springs.len() {
        if springs[0..i].contains(&Condition::Bad) || i + groups[0] > springs.len() {
            break;
        }

        if springs[i..i + groups[0]].contains(&Condition::Good) {
            continue;
        }

        if groups.len() == 1 {
            if i + groups[0] == springs.len() {
                count += 1;
                break;
            } else {
                count += possible_count_cache(&springs[i + groups[0]..], &[], cache);
                continue;
            }
        } else if i + groups[0] + 1 > springs.len() {
            break;
        } else if springs[i + groups[0]] == Condition::Bad {
            continue;
        }

        count += possible_count_cache(
            &springs[i + groups[0] + 1..],
            &groups[1..],
            cache,
        );
    }

    cache.insert((springs.len(), groups.len()), count);

    count
}

fn possible_count(springs: &[Condition], groups: &[usize]) -> usize {
    possible_count_cache(springs, groups, &mut HashMap::new())
}

#[aoc(day12, part1)]
fn part1(input: &Vec<(Vec<Condition>, Vec<usize>)>) -> usize {
    input.iter()
        .map(|(springs, groups)| possible_count(springs, groups))
        .sum()
}

#[aoc(day12, part2)]
fn part2(input: &Vec<(Vec<Condition>, Vec<usize>)>) -> usize {
    input.iter()
        .map(|(springs, groups)| possible_count(
            &[
                &springs[..],
                slice::from_ref(&Condition::WhoKnows),
                &springs[..],
                slice::from_ref(&Condition::WhoKnows),
                &springs[..],
                slice::from_ref(&Condition::WhoKnows),
                &springs[..],
                slice::from_ref(&Condition::WhoKnows),
                &springs[..],
            ].concat(),
            &[
                &groups[..],
                &groups[..],
                &groups[..],
                &groups[..],
                &groups[..],
            ].concat()))
        .sum()
}

#[derive(Debug, PartialEq, Clone)]
enum Condition {
    Good,
    Bad,
    WhoKnows
}

trait ConditionConvertable {
    fn to_condition(self) -> Condition;
}

impl ConditionConvertable for char {
    fn to_condition(self) -> Condition {
        match self {
            '.' => Condition::Good,
            '#' => Condition::Bad,
            '?' => Condition::WhoKnows,
            _ => panic!("Invalid spring char AHHH")
        }
    }
}


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

    const EX: &str = r"???.### 1,1,3
.??..??...?##. 1,1,3
?#?#?#?#?#?#?#? 1,3,1,6
????.#...#... 4,1,1
????.######..#####. 1,6,5
?###???????? 3,2,1";

    #[test]
    fn part1_example() {
        assert_eq!(part1(&parse(EX)), 21);
    }

    #[test]
    fn part2_example() {
        assert_eq!(part2(&parse(EX)), 525152);
    }
}