use aoc_runner_derive::{aoc, aoc_generator};
use nom::{
    bytes::complete::tag,
    character::complete::{newline, space1, u64 as parse_u64},
    error::Error,
    multi::separated_list1,
    sequence::separated_pair,
};

type Input = Vec<(u64, Vec<u64>)>;

#[aoc_generator(day7)]
fn parse(input: &str) -> Input {
    separated_list1(
        newline::<&str, Error<&str>>,
        separated_pair(parse_u64, tag(": "), separated_list1(space1, parse_u64)),
    )(input)
    .unwrap()
    .1
}

#[aoc(day7, part1)]
fn part1(input: &Input) -> usize {
    input
        .iter()
        .filter(|v| part1_test(v.1[0] as i64, v.0 as i64, &v.1[1..]))
        .fold(0, |acc, v| acc + v.0 as usize)
}

fn part1_test(target: i64, remaining: i64, numbers: &[u64]) -> bool {
    if numbers.is_empty() {
        remaining == target
    } else {
        let num = numbers[numbers.len() - 1] as i64;
        let rest = &numbers[0..numbers.len() - 1];
        let mul = remaining % num == 0 && part1_test(target, remaining / num, rest);
        let add = part1_test(target, remaining - num, rest);
        mul || add
    }
}

#[aoc(day7, part2)]
fn part2(input: &Input) -> usize {
    input
        .iter()
        .filter(|v| part2_test(v.1[0] as i64, v.0 as i64, &v.1[1..]))
        .fold(0, |acc, v| acc + v.0 as usize)
}

fn part2_test(target: i64, remaining: i64, numbers: &[u64]) -> bool {
    if numbers.is_empty() {
        remaining == target
    } else {
        let num = numbers[numbers.len() - 1] as i64;
        let rest = &numbers[0..numbers.len() - 1];
        let mul = remaining % num == 0 && part2_test(target, remaining / num, rest);
        let pow = 10i64.pow(num.ilog10() + 1);
        let con = remaining % pow == num && part2_test(target, remaining / pow, rest);
        let add = part2_test(target, remaining - num, rest);
        mul || con || add
    }
}

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

    #[test]
    fn part1_example() {
        assert_eq!(part1(&parse("190: 10 19\n3267: 81 40 27\n83: 17 5\n156: 15 6\n7290: 6 8 6 15\n161011: 16 10 13\n192: 17 8 14\n21037: 9 7 18 13\n292: 11 6 16 20")), 3749);
        assert_eq!(
            part1(&parse(include_str!("../input/2024/day7.txt"))),
            3598800864292
        );
    }

    #[test]
    fn part2_example() {
        assert_eq!(part2(&parse("190: 10 19\n3267: 81 40 27\n83: 17 5\n156: 15 6\n7290: 6 8 6 15\n161011: 16 10 13\n192: 17 8 14\n21037: 9 7 18 13\n292: 11 6 16 20")), 11387);
        assert_eq!(
            part2(&parse(include_str!("../input/2024/day7.txt"))),
            340362529351427
        );
    }
}