use crate::prelude::*;

pub fn solution() -> Solution {
  Solution::new(Day::new(9, 2021), part_1, part_2).with_expected(600, 987840)
}

type Coordinate = (isize, isize);

type HeightMap = HashMap<Coordinate, isize>;

const ADJACENT_OFFSETS: &[Coordinate] = &[
  (-1, 0), // Left
  (1, 0),  // Right
  (0, -1), // Up
  (0, 1),  // Down
];

fn parse_heightmap(input: &str) -> Result<HeightMap> {
  let mut height_map = HeightMap::new();

  for (y, line) in input.lines().enumerate() {
    for (x, height) in line.char_indices() {
      let height = height
        .to_digit(10)
        .ok_or_else(|| eyre!("Invalid input: {}", line))?
        .try_into()?;

      height_map.insert((x.try_into()?, y.try_into()?), height);
    }
  }

  Ok(height_map)
}

fn parse_map_bounds(input: &str) -> Result<Coordinate> {
  Ok((
    input
      .lines()
      .next()
      .map(|line| line.chars().count())
      .ok_or_else(|| eyre!("Invalid input: {}", input))?
      .try_into()?,
    input.lines().count().try_into()?,
  ))
}

fn discover_basins(
  coordinate: Coordinate,
  discovered_coordinates: &mut Vec<Coordinate>,
  height_map: &HeightMap,
) -> Vec<Coordinate> {
  let mut coordinates = vec![];

  for (x_offset, y_offset) in ADJACENT_OFFSETS {
    let neighbor = (x_offset + coordinate.0, y_offset + coordinate.1);
    if discovered_coordinates.contains(&neighbor) {
      continue;
    }

    discovered_coordinates.push(neighbor);
    match height_map.get(&neighbor) {
      Some(9) | None => continue,
      Some(_) => {
        coordinates.push(neighbor);
        coordinates.append(&mut discover_basins(
          neighbor,
          discovered_coordinates,
          height_map,
        ));
      }
    }
  }

  coordinates
}

fn part_1(input: &str) -> Result<String> {
  let height_map = parse_heightmap(input)?;
  let map_bounds = parse_map_bounds(input)?;
  let mut risk_level = 0;

  for y in 0..map_bounds.1 {
    'x_loop: for x in 0..map_bounds.0 {
      let coordinate = (x, y);
      let height = height_map
        .get(&coordinate)
        .ok_or_else(|| eyre!("Coordinate {:?} not found", coordinate))?;

      for (x_offset, y_offset) in ADJACENT_OFFSETS {
        let neighbor = (x_offset + x, y_offset + y);

        if let Some(neighbor_height) = height_map.get(&neighbor) {
          if neighbor_height <= height {
            continue 'x_loop;
          }
        }
      }

      risk_level += 1 + height;
    }
  }

  Ok(risk_level.to_string())
}

fn part_2(input: &str) -> Result<String> {
  let height_map = parse_heightmap(input)?;
  let map_bounds = parse_map_bounds(input)?;
  let mut basins = vec![];
  let mut discovered_coordinates = vec![];
  for y in 0..map_bounds.1 {
    'x_loop: for x in 0..map_bounds.0 {
      let coordinate = (x, y);
      if discovered_coordinates.contains(&coordinate) {
        continue 'x_loop;
      }

      let height = height_map
        .get(&coordinate)
        .ok_or_else(|| eyre!("Coordinate {:?} not found", coordinate))?;
      if height == &9 {
        continue 'x_loop;
      }

      let mut basin = vec![coordinate];
      discovered_coordinates.push(coordinate);
      basin.append(&mut discover_basins(
        coordinate,
        &mut discovered_coordinates,
        &height_map,
      ));
      basins.push(basin.len());
      discovered_coordinates.append(&mut basin);
    }
  }

  Ok(
    basins
      .into_iter()
      .sorted_by(|a, b| b.cmp(a))
      .take(3)
      .product::<usize>()
      .to_string(),
  )
}