use itertools::Itertools;

#[derive(Debug, Clone, Copy, Default, Hash, PartialEq, Eq)]
struct TrailTile {
    coords: (i32, i32),
    height: u8,
    up: Option<bool>,
    down: Option<bool>,
    left: Option<bool>,
    right: Option<bool>,
}

impl TrailTile {
    fn get_neighbors(&self, map: &[Vec<TrailTile>]) -> Vec<TrailTile> {
        let mut neighbors = Vec::new();
        if let Some(up) = self.up {
            if up {
                let neighbor = map[self.coords.1 as usize - 1][self.coords.0 as usize];
                if neighbor.height == self.height + 1 {
                    neighbors.push(neighbor);
                }
            }
        }
        if let Some(down) = self.down {
            if down {
                let neighbor = map[self.coords.1 as usize + 1][self.coords.0 as usize];
                if neighbor.height == self.height + 1 {
                    neighbors.push(neighbor);
                }
            }
        }
        if let Some(left) = self.left {
            if left {
                let neighbor = map[self.coords.1 as usize][self.coords.0 as usize - 1];
                if neighbor.height == self.height + 1 {
                    neighbors.push(neighbor);
                }
            }
        }
        if let Some(right) = self.right {
            if right {
                let neighbor = map[self.coords.1 as usize][self.coords.0 as usize + 1];
                if neighbor.height == self.height + 1 {
                    neighbors.push(neighbor);
                }
            }
        }
        neighbors
    }
}

pub fn process_part1(input: &str) -> i32 {
    let map = input
        .lines()
        .enumerate()
        .map(|(yidx, line)| {
            line.as_bytes()
                .iter()
                .enumerate()
                .map(|(xidx, &height)| TrailTile {
                    coords: (xidx as i32, yidx as i32),
                    height,
                    ..Default::default()
                })
                .collect_vec()
        })
        .collect_vec();
    let mut updated_map = Vec::new();
    let mut trailheads = Vec::new();
    for row in map.iter() {
        let mut new_row = Vec::new();
        for tile in row {
            let up = match map.get((tile.coords.1 - 1) as usize) {
                Some(row) => {
                    let up = row[tile.coords.0 as usize];
                    Some(up.height == tile.height + 1)
                }
                None => None,
            };
            let down = match map.get((tile.coords.1 + 1) as usize) {
                Some(row) => {
                    let down = row[tile.coords.0 as usize];
                    Some(down.height == tile.height + 1)
                }
                None => None,
            };
            let left = map[tile.coords.1 as usize]
                .get((tile.coords.0 - 1) as usize)
                .map(|left| left.height == tile.height + 1);
            let right = map[tile.coords.1 as usize]
                .get((tile.coords.0 + 1) as usize)
                .map(|right| right.height == tile.height + 1);
            let tile = TrailTile {
                up,
                down,
                left,
                right,
                ..*tile
            };
            if tile.height == b'0' {
                trailheads.push(tile);
            }
            //print!(
            //    "{}: {}, {}      ",
            //    tile.height as char, tile.coords.0, tile.coords.1
            //);
            new_row.push(tile);
        }
        //println!();
        updated_map.push(new_row);
    }
    let mut res = 0;
    //println!("0 {trailheads:?}");
    for trailhead in trailheads {
        let mut neighbors = trailhead.get_neighbors(&updated_map);
        //println!("1 {neighbors:?}");
        for _idx in 2..=9 {
            //println!("{idx}");
            if neighbors.is_empty() {
                //println!("empty");
                break;
            }
            let inner_neighbors = neighbors.clone();
            neighbors.clear();
            for neighbor in inner_neighbors {
                neighbors.extend_from_slice(&neighbor.get_neighbors(&updated_map));
            }
            //println!("{idx} {neighbors:?}");
        }
        res += neighbors.iter().unique().count();
    }
    res as i32
}

pub fn process_part2(input: &str) -> i32 {
    let map = input
        .lines()
        .enumerate()
        .map(|(yidx, line)| {
            line.as_bytes()
                .iter()
                .enumerate()
                .map(|(xidx, &height)| TrailTile {
                    coords: (xidx as i32, yidx as i32),
                    height,
                    ..Default::default()
                })
                .collect_vec()
        })
        .collect_vec();
    let mut updated_map = Vec::new();
    let mut trailheads = Vec::new();
    for row in map.iter() {
        let mut new_row = Vec::new();
        for tile in row {
            let up = match map.get((tile.coords.1 - 1) as usize) {
                Some(row) => {
                    let up = row[tile.coords.0 as usize];
                    Some(up.height == tile.height + 1)
                }
                None => None,
            };
            let down = match map.get((tile.coords.1 + 1) as usize) {
                Some(row) => {
                    let down = row[tile.coords.0 as usize];
                    Some(down.height == tile.height + 1)
                }
                None => None,
            };
            let left = map[tile.coords.1 as usize]
                .get((tile.coords.0 - 1) as usize)
                .map(|left| left.height == tile.height + 1);
            let right = map[tile.coords.1 as usize]
                .get((tile.coords.0 + 1) as usize)
                .map(|right| right.height == tile.height + 1);
            let tile = TrailTile {
                up,
                down,
                left,
                right,
                ..*tile
            };
            if tile.height == b'0' {
                trailheads.push(tile);
            }
            //print!(
            //    "{}: {}, {}      ",
            //    tile.height as char, tile.coords.0, tile.coords.1
            //);
            new_row.push(tile);
        }
        //println!();
        updated_map.push(new_row);
    }
    let mut res = 0;
    //println!("0 {trailheads:?}");
    for trailhead in trailheads {
        let mut neighbors = trailhead.get_neighbors(&updated_map);
        //println!("1 {neighbors:?}");
        for _idx in 2..=9 {
            //println!("{idx}");
            if neighbors.is_empty() {
                //println!("empty");
                break;
            }
            let inner_neighbors = neighbors.clone();
            neighbors.clear();
            for neighbor in inner_neighbors {
                neighbors.extend_from_slice(&neighbor.get_neighbors(&updated_map));
            }
            //println!("{idx} {neighbors:?}");
        }
        res += neighbors.len();
    }
    res as i32
}

#[cfg(test)]
mod tests {
    use super::*;

    const INPUT_SMALL: &str = "0123
1234
8765
9876";

    const INPUT_LARGE: &str = "89010123
78121874
87430965
96549874
45678903
32019012
01329801
10456732";

    const INPUT_SIMPLIFIED_1: &str = "...0...
...1...
...2...
6543456
7.....7
8.....8
9.....9";
    const INPUT_SIMPLIFIED_2: &str = "..90..9
...1.98
...2..7
6543456
765.987
876....
987....";
    const INPUT_SIMPLIFIED_3: &str = "10..9..
2...8..
3...7..
4567654
...8..3
...9..2
.....01";

    const INPUT_SIMPLIFIED_4: &str = ".....0.
..4321.
..5..2.
..6543.
..7..4.
..8765.
..9....";

    const INPUT_SIMPLIFIED_5: &str = "012345
123456
234567
345678
4.6789
56789.";

    #[test]
    fn part1_small() {
        let result = process_part1(INPUT_SMALL);
        assert_eq!(result, 1);
    }

    #[test]
    fn part1_large() {
        let result = process_part1(INPUT_LARGE);
        assert_eq!(result, 36);
    }

    #[test]
    fn part1_simplified_1() {
        let result = process_part1(INPUT_SIMPLIFIED_1);
        assert_eq!(result, 2);
    }

    #[test]
    fn part1_simplified_2() {
        let result = process_part1(INPUT_SIMPLIFIED_2);
        assert_eq!(result, 4);
    }

    #[test]
    fn part1_simplified_3() {
        let result = process_part1(INPUT_SIMPLIFIED_3);
        assert_eq!(result, 3);
    }

    #[test]
    fn part2_small() {
        let result = process_part2(INPUT_SMALL);
        assert_eq!(result, 16);
    }

    #[test]
    fn part2_large() {
        let result = process_part2(INPUT_LARGE);
        assert_eq!(result, 81);
    }

    #[test]
    fn part2_simplified_2() {
        let result = process_part2(INPUT_SIMPLIFIED_2);
        assert_eq!(result, 13);
    }

    #[test]
    fn part2_simplified_4() {
        let result = process_part2(INPUT_SIMPLIFIED_4);
        assert_eq!(result, 3);
    }

    #[test]
    fn part2_simplified_5() {
        let result = process_part2(INPUT_SIMPLIFIED_5);
        assert_eq!(result, 227);
    }
}