use nom::{
    branch::alt,
    bytes::complete::tag,
    character::complete::line_ending,
    combinator::{map, value},
    multi::{many1, separated_list1},
    IResult,
};
use std::fs;

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let input = fs::read_to_string("inputs/day_25.txt")?;
    let parsed = SnafuList::parser(&input).unwrap().1;
    dbg!(to_snafu(parsed.sum()));

    Ok(())
}

#[derive(Debug, Clone, PartialEq, Eq)]
struct SnafuList(Vec<i64>);

impl SnafuList {
    fn parser(input: &str) -> IResult<&str, SnafuList> {
        map(
            separated_list1(
                line_ending,
                map(
                    many1(alt((
                        value(2, tag("2")),
                        value(1, tag("1")),
                        value(0, tag("0")),
                        value(-1, tag("-")),
                        value(-2, tag("=")),
                    ))),
                    |digits| {
                        let mut result: i64 = 0;
                        for digit in digits {
                            result *= 5;
                            result += digit;
                        }
                        result
                    },
                ),
            ),
            SnafuList,
        )(input)
    }

    fn sum(&self) -> i64 {
        self.0.iter().sum()
    }
}

fn to_snafu(mut remaining_value: i64) -> String {
    let mut result = String::new();

    let mut current_power = 5i64.pow(26);

    while current_power > 0 {
        let next_digit_range = {
            let mut next_digit_power = current_power / 5;
            let mut max_next_digit = 0;
            while next_digit_power > 0 {
                max_next_digit += 2 * next_digit_power;
                next_digit_power /= 5;
            }
            -max_next_digit..=max_next_digit
        };

        let (digit, digit_value) = [('=', -2), ('-', -1), ('0', 0), ('1', 1), ('2', 2)]
            .into_iter()
            .filter_map(|(digit, digit_value)| {
                let digit_value = digit_value * current_power;
                let remaining_if_digit_set = remaining_value - digit_value;
                if next_digit_range.contains(&remaining_if_digit_set) {
                    Some((digit, digit_value))
                } else {
                    None
                }
            })
            .next()
            .expect("No digit found");

        remaining_value -= digit_value;
        result.push(digit);

        current_power /= 5;
    }

    result.trim_start_matches("0").to_owned()
}

#[test]
fn round_trip_works() {
    let str_list = vec![
        "1=-0-2", "12111", "2=0=", "21", "2=01", "111", "20012", "112", "1=-1=", "1-12", "12",
        "1=", "122",
    ];
    let list = SnafuList::parser(&str_list.join("\n")).unwrap().1;

    assert_eq!(
        list,
        SnafuList(vec![
            1747, 906, 198, 11, 201, 31, 1257, 32, 353, 107, 7, 3, 37
        ])
    );

    for (i, num) in list.0.iter().enumerate() {
        assert_eq!(to_snafu(*num), str_list[i]);
    }
}