use nom::{
    bytes::complete::tag,
    character::complete::{alpha1, anychar, line_ending},
    combinator::map,
    multi::{many0, separated_list1},
    sequence::tuple,
    IResult,
};
use std::{collections::BTreeMap, fs};

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let input = fs::read_to_string("inputs/day_14.txt")?;
    let mut polymer = parse_polymer_expansion(&input).unwrap().1;
    for _ in 0..10 {
        polymer.expand();
    }
    dbg!(polymer.most_common_element_count() - polymer.least_common_element_count());
    for _ in 0..30 {
        polymer.expand();
    }
    dbg!(polymer.most_common_element_count() - polymer.least_common_element_count());

    Ok(())
}

#[derive(Debug)]
struct PolymerExpansion {
    polymer: Vec<char>,
    rules: Rules,
    element_counts: BTreeMap<char, u64>,
}

impl PolymerExpansion {
    fn new(polymer: Vec<char>, rules: Rules) -> PolymerExpansion {
        let mut element_counts = BTreeMap::new();
        for c in &polymer {
            *element_counts.entry(*c).or_insert(0) += 1;
        }

        PolymerExpansion {
            polymer,
            rules,
            element_counts,
        }
    }

    fn expand(&mut self) {
        let mut next_polymer = Vec::new();
        for i in 0..self.polymer.len() - 1 {
            next_polymer.push(self.polymer[i]);
            if let Some(extra) = self.rules.get(&[self.polymer[i], self.polymer[i + 1]]) {
                next_polymer.push(extra);
                *self.element_counts.entry(extra).or_insert(0) += 1;
            }
        }
        next_polymer.push(self.polymer[self.polymer.len() - 1]);
        self.polymer = next_polymer;
    }

    fn most_common_element_count(&self) -> u64 {
        self.element_counts.values().max().cloned().unwrap_or(0)
    }

    fn least_common_element_count(&self) -> u64 {
        self.element_counts.values().min().cloned().unwrap_or(0)
    }
}

#[derive(Debug)]
struct Rules {
    rules: BTreeMap<[char; 2], char>,
}

impl Rules {
    fn get(&self, key: &[char; 2]) -> Option<char> {
        self.rules.get(key).cloned()
    }
}

fn parse_polymer_expansion(input: &str) -> IResult<&str, PolymerExpansion> {
    let (input, polymer) = map(alpha1, |s: &str| s.chars().collect())(input)?;
    let (input, _) = many0(line_ending)(input)?;
    let (input, rules) = parse_rules(input)?;
    Ok((input, PolymerExpansion::new(polymer, rules)))
}

fn parse_rules(input: &str) -> IResult<&str, Rules> {
    map(separated_list1(line_ending, parse_rule), |rules| Rules {
        rules: rules.into_iter().collect(),
    })(input)
}

fn parse_rule(input: &str) -> IResult<&str, ([char; 2], char)> {
    map(
        tuple((anychar, anychar, tag(" -> "), anychar)),
        |(p1, p2, _, r)| ([p1, p2], r),
    )(input)
}