use std::collections::HashSet;

use color_eyre::eyre;
use eyre::eyre;

use crate::lib::{IterExt, Vec2};


#[derive(Debug)]
struct Cell {
    height: u32,
    min: bool,
}

struct NeighborsIter {
    x: usize,
    y: usize,
    rows: usize,
    cols: usize,
    state: usize,
}

impl Iterator for NeighborsIter {
    type Item = (usize, usize);
    fn next(&mut self) -> Option<Self::Item> {
        let neighbor = match self.state {
            0 if self.x > 0 => Some((self.y, self.x - 1)),
            1 if self.y > 0 => Some((self.y - 1, self.x)),
            2 if self.x < self.cols - 1 => Some((self.y, self.x + 1)),
            3 if self.y < self.rows - 1 => Some((self.y + 1, self.x)),
            _ => None,
        };
        self.state += 1;

        if self.state > 4 {
            None
        }
        else {
            neighbor.or_else(|| self.next())
        }
    }
}


// struct CellWindow {
//     before: Option<Cell>,
//     cell: Cell,
//     above: Option<Cell>,
// }


fn load(data: &str) -> eyre::Result<Vec2<Cell>> {
    let num_columns = data.lines().next().unwrap().len();
    let mut heights = Vec2::new(num_columns);
    for line in data.lines() {
        let mut row = Vec::new();
        for c in line.chars() {
            match c.to_digit(10) {
                Some(d) => row.push(Cell {height: d, min: false}),
                None => return Err(eyre!("Invalid character (not a digit): {}", c)),
            }
        }
        heights.push(row);
    }
    Ok(heights)
}


fn scan(map: &mut Vec2<Cell>) {
    for row in 0..map.row_count() {
        for col in 0..map.col_count() {
            let mut lt_left = true;
            if col > 0 {
                if map[row][col].height < map[row][col - 1].height {
                    map[row][col - 1].min = false;
                }
                else {
                    lt_left = false;
                }
            }

            let mut lt_top = true;
            if row > 0 {
                if map[row][col].height < map[row - 1][col].height {
                    map[row - 1][col].min = false;
                }
                else {
                    lt_top = false;
                }
            }

            if lt_left && lt_top {
                map[row][col].min = true;
            }
        }
    }
}


fn low_points(map: &Vec2<Cell>) -> impl Iterator<Item=(usize, &Cell)> {
    map.values().enumerate().filter(|(_, c)| c.min)
}


fn neighbors(map: &Vec2<Cell>, coords: (usize, usize)) -> NeighborsIter {
    NeighborsIter {
        x: coords.1,
        y: coords.0,
        rows: map.row_count(),
        cols: map.col_count(),
        state: 0,
    }
}


fn basin_size (map: &Vec2<Cell>, start: (usize, usize), visited: &mut HashSet<(usize, usize)>) -> u32 {
    let mut total = 1;
    for coords in neighbors(map, start) {
        if map[coords.0][coords.1].height < 9 && !visited.contains(&coords) {
            visited.insert(coords);
            total += basin_size(map, coords, visited);
        }
    }
    total
}


fn basins(map: &Vec2<Cell>) -> u32 {
    low_points(map)
        .map(|(i, _)| {
            let start = map.coords(i);
            basin_size(map, start, &mut HashSet::from([start]))
        })
        .max_n(3)
        .unwrap()
        .iter()
        .product()
}


pub fn run(data: &str) -> eyre::Result<(u32, u32)> {
    //  let sample = "2199943210\n3987894921\n9856789892\n8767896789\n9899965678";
    let mut map = load(data)?;
    scan(&mut map);

    let one = low_points(&map).fold(0, |acc, (_, c)| acc + c.height + 1);
    let two = basins(&map);

    Ok((one, two))
}