use nom::{
    character::complete::{line_ending, one_of},
    combinator::{map, map_res},
    multi::{many1, separated_list1},
    IResult,
};
use std::{collections::BTreeMap, fs};

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let input = fs::read_to_string("inputs/day_15.txt")?;
    let risk_grid = parse_risk_grid(&input).unwrap().1;
    let big_risk_grid = risk_grid.expand();
    {
        let mut searcher = Searcher::new(risk_grid);
        while searcher.end_risk().is_none() {
            searcher.explore_next();
        }
        dbg!(searcher.end_risk());
    }
    {
        let mut big_searcher = Searcher::new(big_risk_grid);
        while big_searcher.end_risk().is_none() {
            big_searcher.explore_next();
        }
        dbg!(big_searcher.end_risk());
    }

    Ok(())
}

#[derive(Debug)]
struct Searcher {
    grid: RiskGrid,
    frontier: Vec<(Risk, Point)>,
    explored: BTreeMap<Point, Risk>,
}

impl Searcher {
    fn new(grid: RiskGrid) -> Searcher {
        let start = grid.start();
        let mut explored = BTreeMap::new();
        explored.insert(start.clone(), Risk(0));
        Searcher {
            grid,
            explored,
            frontier: vec![(Risk(0), start)],
        }
    }

    fn explore_next(&mut self) {
        if let Some((next_risk, next_point)) = self.frontier.pop() {
            for (neighbour_risk, neighbour_point) in self.grid.neighbours(&next_point) {
                if !self.explored.contains_key(&neighbour_point) {
                    let total_risk = next_risk + neighbour_risk;
                    self.explored.insert(neighbour_point.clone(), total_risk);
                    self.frontier.push((total_risk, neighbour_point));
                }
            }
            self.frontier.sort_unstable_by(|a, b| b.0.cmp(&a.0));
        }
    }

    fn end_risk(&self) -> Option<Risk> {
        self.explored.get(&self.grid.end()).cloned()
    }
}

#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
struct Point {
    x: usize,
    y: usize,
}

#[derive(Debug)]
struct RiskGrid {
    map: Vec<Vec<Risk>>,
}

impl RiskGrid {
    fn start(&self) -> Point {
        Point { x: 0, y: 0 }
    }

    fn end(&self) -> Point {
        let y = self.map.len() - 1;
        let x = self.map[y].len() - 1;
        Point { x, y }
    }

    fn risk_at_point(&self, p: &Point) -> Option<Risk> {
        self.map.get(p.y).and_then(|row| row.get(p.x)).cloned()
    }

    fn neighbours(&self, p: &Point) -> Vec<(Risk, Point)> {
        let mut neighbours = Vec::new();
        if p.x > 0 {
            let left = Point { x: p.x - 1, y: p.y };
            if let Some(risk) = self.risk_at_point(&left) {
                neighbours.push((risk, left));
            }
        }
        if p.y > 0 {
            let up = Point { x: p.x, y: p.y - 1 };
            if let Some(risk) = self.risk_at_point(&up) {
                neighbours.push((risk, up));
            }
        }
        let right = Point { x: p.x + 1, y: p.y };
        if let Some(risk) = self.risk_at_point(&right) {
            neighbours.push((risk, right));
        }
        let down = Point { x: p.x, y: p.y + 1 };
        if let Some(risk) = self.risk_at_point(&down) {
            neighbours.push((risk, down));
        }
        neighbours
    }

    fn expand(&self) -> RiskGrid {
        let mut new_map = Vec::new();
        for y_repeat in 0..5 {
            for original_row in &self.map {
                let mut new_row = Vec::new();
                for x_repeat in 0..5 {
                    let risk_modifier = y_repeat + x_repeat;
                    for original_risk in original_row {
                        let modified_risk = original_risk.grid_wrap_increment(risk_modifier);
                        new_row.push(modified_risk);
                    }
                }
                new_map.push(new_row);
            }
        }
        RiskGrid { map: new_map }
    }
}

#[derive(Debug, Clone, Copy, derive_more::Add, PartialEq, Eq, PartialOrd, Ord)]
struct Risk(u64);

impl Risk {
    fn grid_wrap_increment(&self, increment_count: u64) -> Risk {
        let new = (self.0 + increment_count) % 9;
        if new == 0 {
            Risk(9)
        } else {
            Risk(new)
        }
    }
}

fn parse_risk_grid(input: &str) -> IResult<&str, RiskGrid> {
    map(separated_list1(line_ending, many1(parse_risk)), |risks| {
        RiskGrid { map: risks }
    })(input)
}

fn parse_risk(input: &str) -> IResult<&str, Risk> {
    map_res(one_of("0123456789"), |digit| {
        digit.to_string().parse().map(Risk)
    })(input)
}