use std::collections::HashMap;

use color_eyre::{eyre::eyre, Result};
use itertools::Itertools;

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

type Coordinate = (isize, isize);

type HeightMap = HashMap<Coordinate, isize>;

const ADJACENT_OFFSETS: &[Coordinate] = &[
  (-1, 00), // Left
  (01, 00), // Right
  (00, -1), // Up
  (00, 01), // 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(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(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<isize> {
  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(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)
}

fn part_2(input: &str) -> Result<usize> {
  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(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(),
  )
}

#[test]
fn test_sample() -> Result<()> {
  let sample = "2199943210\n3987894921\n9856789892\n8767896789\n9899965678";
  assert_eq!(part_1(sample)?, 15);
  assert_eq!(part_2(sample)?, 1134);
  Ok(())
}