use crate::prelude::*;

pub fn solution() -> Solution {
  Solution::new(Day::new(10, 2020), part_1, part_2)
    .with_expected(2368, 1727094849536_i64)
}

fn parse_numbers(input: &str) -> (HashMap<usize, i32>, HashMap<i32, u32>) {
  let mut numbers = input
    .lines()
    .map(|line| line.parse().unwrap())
    .collect::<Vec<usize>>();

  numbers.push(0);
  numbers.sort_unstable();
  numbers.push(numbers.last().unwrap() + 3);

  let mut intervals = HashMap::new();
  let mut streaks = HashMap::new();
  let mut streak_length = 0;
  for (index, a) in numbers.iter().enumerate() {
    if let Some(b) = numbers.get(index + 1) {
      let interval = b - a;
      if interval == 1 {
        streak_length += 1;
      } else {
        *streaks.entry(streak_length).or_insert(0) += 1;
        streak_length = 0;
      }

      *intervals.entry(interval).or_insert(0) += 1;
    }
  }

  (intervals, streaks)
}

fn part_1(input: &str) -> Result<String> {
  let (intervals, _) = parse_numbers(input);
  Ok((intervals.get(&1).unwrap() * intervals.get(&3).unwrap()).to_string())
}

fn part_2(input: &str) -> Result<String> {
  let (_, streaks) = parse_numbers(input);
  Ok(
    (7_usize.pow(*streaks.get(&4).unwrap())
      * 4_usize.pow(*streaks.get(&3).unwrap())
      * 2_usize.pow(*streaks.get(&2).unwrap()))
    .to_string(),
  )
}