use crate::prelude::*;

pub fn solution() -> Solution {
  Solution::new(Day::new(3, 2021), part_1, part_2)
    .with_expected(1092896, 4672151)
}

fn count_bits(input: &str) -> Result<Vec<(usize, i32)>> {
  let mut bits = HashMap::<usize, i32>::new();

  for line in input.lines() {
    for (index, bit) in line.char_indices() {
      let entry = bits.entry(index).or_default();

      match bit {
        '0' => *entry -= 1,
        '1' => *entry += 1,
        _ => return Err(eyre!("Unknown character in line: {}", line)),
      }
    }
  }

  Ok(
    bits
      .into_iter()
      .sorted_by(|(a, _), (b, _)| a.cmp(b))
      .collect(),
  )
}

fn part_1(input: &str) -> Result<String> {
  let bits = count_bits(input)?;

  let gamma_rate = i32::from_str_radix(
    &bits
      .clone()
      .into_iter()
      .map(|(_, bit)| if bit > 0 { '1' } else { '0' })
      .collect::<String>(),
    2,
  )?;

  let epsilon_rate = i32::from_str_radix(
    &bits
      .into_iter()
      .map(|(_, bit)| if bit < 0 { '1' } else { '0' })
      .collect::<String>(),
    2,
  )?;

  Ok((gamma_rate * epsilon_rate).to_string())
}

fn part_2(input: &str) -> Result<String> {
  let mut most_common_lines = input.lines().collect::<Vec<_>>();
  let mut least_common_lines = input.lines().collect::<Vec<_>>();

  let mut most_common_bits = count_bits(input)?;
  let mut least_common_bits = count_bits(input)?;

  for index in 0..most_common_bits.len() {
    if most_common_lines.len() > 1 {
      let (index, bit) = most_common_bits
        .get(index)
        .ok_or_else(|| eyre!("Could not find most common bit"))?;
      let most_common = if bit >= &0 { '1' } else { '0' };
      most_common_lines = most_common_lines
        .into_iter()
        .filter(|line| line.chars().nth(*index) == Some(most_common))
        .collect();
      most_common_bits = count_bits(&most_common_lines.join("\n"))?;
    }

    if least_common_lines.len() > 1 {
      let (index, bit) = least_common_bits
        .get(index)
        .ok_or_else(|| eyre!("Could not find least common bit"))?;
      let least_common = if bit < &0 { '1' } else { '0' };
      least_common_lines = least_common_lines
        .into_iter()
        .filter(|line| line.chars().nth(*index) == Some(least_common))
        .collect();
      least_common_bits = count_bits(&least_common_lines.join("\n"))?;
    }
  }

  let oxygen_generator_rating = i32::from_str_radix(
    most_common_lines
      .first()
      .ok_or_else(|| eyre!("Didn't find an oxygen generator rating"))?,
    2,
  )?;
  let co2_scrubber_rating = i32::from_str_radix(
    least_common_lines
      .first()
      .ok_or_else(|| eyre!("Didn't find a CO2 scrubber rating"))?,
    2,
  )?;

  Ok((oxygen_generator_rating * co2_scrubber_rating).to_string())
}