use std::collections::HashMap;

use color_eyre::Result;

type FishMap = HashMap<isize, isize>;

pub fn solve() -> Result<()> {
  let input_data = include_str!("../../data/day_06.txt").trim();
  println!("Day 06 Part 1: {}", part_1(input_data)?);
  println!("Day 06 Part 2: {}", part_2(input_data)?);
  Ok(())
}

fn part_1(input: &str) -> Result<isize> {
  let mut fishes = parse_fishes(input)?;

  for _ in 0..80 {
    fishes = simulate_fishes(fishes);
  }

  Ok(count_fishes(fishes))
}

fn part_2(input: &str) -> Result<isize> {
  let mut fishes = parse_fishes(input)?;

  for _ in 0..256 {
    fishes = simulate_fishes(fishes);
  }

  Ok(count_fishes(fishes))
}

fn parse_fishes(input: &str) -> Result<FishMap> {
  let mut fishes = FishMap::new();
  let individual_fishes = input
    .split(",")
    .map(str::parse)
    .collect::<Result<Vec<_>, _>>()?;

  for fish in individual_fishes {
    let amount = fishes.entry(fish).or_default();
    *amount += 1;
  }

  Ok(fishes)
}

fn count_fishes(fishes: FishMap) -> isize {
  fishes.into_iter().map(|(_, amount)| amount).sum()
}

fn simulate_fishes(fishes: FishMap) -> FishMap {
  let mut new_fishes = HashMap::new();

  for (mut timer, amount) in fishes {
    timer -= 1;

    if timer < 0 {
      let new_fish = new_fishes.entry(8).or_default();
      *new_fish += amount;

      let reset_fish = new_fishes.entry(6).or_default();
      *reset_fish += amount;
    } else {
      let current_fish = new_fishes.entry(timer).or_default();
      *current_fish += amount;
    }
  }

  new_fishes
}