use std::collections::HashMap;

use super::Solution;

pub struct Day12 {}

impl Solution for Day12 {
    fn part1(
        &self,
        input: &mut Vec<String>,
    ) -> Result<Box<dyn std::fmt::Display + Sync>, Box<dyn std::error::Error>> {
        let mut ans = 0usize;
        for s in input {
            let springs = s.split_whitespace().nth(0).unwrap().to_owned();
            let nums = s
                .split_whitespace()
                .nth(1)
                .unwrap()
                .split(',')
                .map(|n| n.parse::<usize>().unwrap())
                .collect::<Vec<_>>();

            ans += self.solve(&springs.as_bytes(), None, &nums, &mut HashMap::new());
        }
        Ok(Box::new(ans))
    }

    fn part2(
        &self,
        input: &mut Vec<String>,
    ) -> Result<Box<dyn std::fmt::Display + Sync>, Box<dyn std::error::Error>> {
        for l in input.iter_mut() {
            let (s, n) = l.split_once(" ").unwrap();
            *l = format!("{s}?{s}?{s}?{s}?{s} {n},{n},{n},{n},{n}\n");
        }

        let mut ans = 0usize;
        for s in input {
            let springs = s.split_whitespace().nth(0).unwrap().to_owned();
            let nums = s
                .split_whitespace()
                .nth(1)
                .unwrap()
                .split(',')
                .map(|n| n.parse::<usize>().unwrap())
                .collect::<Vec<_>>();

            ans += self.solve(&springs.as_bytes(), None, &nums, &mut HashMap::new());
        }
        Ok(Box::new(ans))
    }

    fn get_day(&self) -> u8 {
        12
    }
}

impl Day12 {
    /// recursively counts the number of permutations of spring we could get from
    /// the splits specified in 'nums'
    fn solve<'a, 'b>(
        &self,
        s: &'a [u8],
        in_group: Option<usize>,
        cons: &'b [usize],
        map: &mut HashMap<(&'a [u8], Option<usize>, &'b [usize]), usize>,
    ) -> usize {
        if s.is_empty() {
            return match in_group {
                Some(n) if cons == &[n] => 1,
                None if cons.is_empty() => 1,
                _ => 0,
            };
        }

        // Check for a cache hit
        if s[0] == b'?' {
            if let Some(result) = map.get(&(s, in_group, cons)) {
                return *result;
            }
        }

        // Resursively match based on the whether we are in a block and/or we have spaces left to fill
        let ans = match (s[0], in_group, cons) {
            (b'.', None, _) | (b'?', None, []) => self.solve(&s[1..], None, cons, map),
            (b'.' | b'?', Some(n), [e, ..]) if n == *e => {
                self.solve(&s[1..], None, &cons[1..], map)
            }
            (b'#' | b'?', Some(n), [e, ..]) if n < *e => {
                self.solve(&s[1..], Some(n + 1), cons, map)
            }
            (b'#', None, [_, ..]) => self.solve(&s[1..], Some(1), cons, map),
            (b'?', None, _) => {
                self.solve(&s[1..], None, cons, map) + self.solve(&s[1..], Some(1), cons, map)
            }
            _ => 0,
        };

        // Store in cache
        if s[0] == b'?' {
            map.insert((s, in_group, cons), ans);
        }

        ans
    }
}

/// Test from puzzle input
#[cfg(test)]
mod test {
    use super::*;
    use crate::*;

    #[test]
    fn part1() {
        let challenge = day12::Day12 {};

        //Complete the Challenge
        let answer = challenge
            .part1(
                utils::get_input(challenge.get_day(), utils::InputType::Test1)
                    .unwrap()
                    .as_mut(),
            )
            .unwrap()
            .to_string();

        assert_eq!(answer, "21");
    }

    #[test]
    fn part2() {
        let challenge = day12::Day12 {};

        //Complete the Challenge
        let answer = challenge
            .part2(
                utils::get_input(challenge.get_day(), utils::InputType::Test2)
                    .unwrap()
                    .as_mut(),
            )
            .unwrap()
            .to_string();

        assert_eq!(answer, "525152");
    }
}