use std::{
    cmp::Reverse,
    collections::{BinaryHeap, HashMap},
};

use anyhow::Context;

use crate::{Problem, Solution};

pub struct Day17;

impl Problem for Day17 {
    const DAY: u8 = 17;

    const INPUT: &'static str = include_str!("../input/day_17.txt");
}

impl Solution for Day17 {
    type Answer1 = usize;

    type Answer2 = usize;

    fn part_1(input: &str) -> anyhow::Result<Self::Answer1> {
        let grid = parse_grid(input).context("Failed to parse grid")?;

        astar(&grid, (0, 0), (grid.len() - 1, grid[0].len() - 1), 1, 3)
            .context("Failed to find path")

        // 755 ==
    }

    fn part_2(input: &str) -> anyhow::Result<Self::Answer2> {
        let grid = parse_grid(input).context("Failed to parse grid")?;

        astar(&grid, (0, 0), (grid.len() - 1, grid[0].len() - 1), 4, 10)
            .context("Failed to find path")
        // 879 <
        // 881 ==
    }
}

fn parse_grid(input: &str) -> Option<Grid> {
    input
        .lines()
        .map(|v| {
            v.chars()
                .map(|c| c.to_digit(10).map(|n| n as usize))
                .try_collect::<Vec<_>>()
        })
        .try_collect::<Grid>()
}

fn astar(
    grid: &Grid,
    start: Position,
    goal: Position,
    min_step: usize,
    max_step: usize,
) -> Option<usize> {
    let start_node = Node {
        pos: start,
        edge: None,
        len: 0,
        path: vec![],
    };
    let mut queue = BinaryHeap::from([(Reverse(0), start_node)]);
    let mut cache = HashMap::from([]);

    while let Some((Reverse(node_cost), node)) = queue.pop() {
        if node.pos == goal {
            // #[cfg(test)]
            // crate::Printer::from(std::io::stdout().lock())
            //     .with_grid(grid, &rebuild_path(node, &cache, start))
            //     .write()
            //     .unwrap();

            return Some(node_cost);
        }

        'edge_loop: for edge in Edge::ALL {
            let mut next = node.clone();

            next.edge = Some(edge);
            next.len = 0;

            let Some((next, mut next_cost)) = extend_node(next, edge, min_step, max_step, grid)
            else {
                continue 'edge_loop;
            };

            next_cost += node_cost;

            if !cache.get(&next).is_some_and(|c| next_cost >= *c) {
                cache.insert(next.clone(), next_cost);
                queue.push((Reverse(next_cost), next));
            }
        }
    }

    None
}

fn extend_node(
    mut node: Node,
    edge: Edge,
    min: usize,
    max: usize,
    grid: &Grid,
) -> Option<(Node, usize)> {
    let mut cost = 0;

    let len = match node.path.last() {
        Some((e, l)) if *e == edge => {
            node.len = *l;
            node.len + 1
        }
        _ => min,
    };

    while node.len < len {
        node.len += 1;
        if node.len > max {
            // println!("Breaking due to len: {}", node.len);
            return None;
        }

        node.pos = node.edge?.traverse(node.pos)?;
        cost += grid.get(node.pos.0)?.get(node.pos.1)?;
    }

    Some((node, cost))
}

// #[cfg(test)]
// fn rebuild_path(node: Node, cache: &HashMap<Node, usize>, start: Position) -> Vec<Position> {
//     let mut path = vec![node.pos];
//     let mut opt = Some((0, node));
//     while let Some((_, next)) = opt {
//         let mut pos = next.pos;
//         for _ in 0..next.len {
//             path.push(pos);
//             pos = next.edge.unwrap().inverse().traverse(pos).unwrap();
//         }
//         opt = cache.get(&next).cloned();
//     }
//     path.push(start);
//
//     path
// }

type Grid = Vec<Vec<usize>>;
type Position = (usize, usize);

#[derive(Debug, Default, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
struct Node {
    pos: (usize, usize),
    edge: Option<Edge>,
    len: usize,
    path: Vec<(Edge, usize)>,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
enum Edge {
    North,
    East,
    South,
    West,
}

impl Edge {
    const ALL: [Edge; 4] = [Edge::North, Edge::East, Edge::South, Edge::West];

    fn edges(&self) -> &[Edge] {
        match self {
            Edge::North => &[Edge::North, Edge::East, Edge::West],
            Edge::East => &[Edge::North, Edge::East, Edge::South],
            Edge::South => &[Edge::East, Self::South, Edge::West],
            Edge::West => &[Edge::North, Edge::South, Edge::West],
        }
    }

    fn traverse(&self, (row, col): Position) -> Option<Position> {
        Some(match self {
            Edge::North => (row.checked_sub(1)?, col),
            Edge::East => (row, col.checked_add(1)?),
            Edge::South => (row.checked_add(1)?, col),
            Edge::West => (row, col.checked_sub(1)?),
        })
    }

    fn inverse(&self) -> Self {
        match self {
            Edge::North => Edge::South,
            Edge::East => Edge::West,
            Edge::South => Edge::North,
            Edge::West => Edge::East,
        }
    }
}

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

    const INPUT: &str = indoc::indoc! {r#"
        2413432311323
        3215453535623
        3255245654254
        3446585845452
        4546657867536
        1438598798454
        4457876987766
        3637877979653
        4654967986887
        4564679986453
        1224686865563
        2546548887735
        4322674655533
    "#};

    const INPUT2: &str = indoc::indoc! {r#"
        111111111111
        999999999991
        999999999991
        999999999991
        999999999991
    "#};

    #[test]
    fn test_part_1() -> anyhow::Result<()> {
        Ok(assert_eq!(102, Day17::part_1(INPUT)?))
    }

    #[test]
    fn test_part_2_1() -> anyhow::Result<()> {
        Ok(assert_eq!(94, Day17::part_2(INPUT)?))
    }

    #[test]
    fn test_part_2_2() -> anyhow::Result<()> {
        Ok(assert_eq!(71, Day17::part_2(INPUT2)?))
    }
}