use std::collections::{/* HashMap, */ HashSet};

use crate::{Problem, Solution};

pub struct Day16;

impl Problem for Day16 {
    const DAY: u8 = 16;

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

impl Solution for Day16 {
    type Answer1 = usize;

    type Answer2 = usize;

    fn part_1(input: &str) -> anyhow::Result<Self::Answer1> {
        let grid = parse_grid(input)?;

        // let mut cache = HashMap::default();
        let score = solve_grid(
            &grid,
            /* &mut cache, */ Default::default(),
            Default::default(),
        );
        let score = score.into_iter().map(|(_, p)| p).collect::<HashSet<_>>();

        // let mut output = vec![vec!["[00]".to_string(); grid[0].len()]; grid.len()];
        //
        // for ((e, (row, col)), score) in cache {
        //     output[col][row] = format!("[{:02}]", score.len());
        //     // println!("{beam:?}: {score}");
        // }
        // println!(
        //     "{}",
        //     output
        //         .into_iter()
        //         .map(|v| v.into_iter().collect::<String>())
        //         .collect::<Vec<_>>()
        //         .join("\n")
        // );

        Ok(score.len())
    }

    fn part_2(input: &str) -> anyhow::Result<Self::Answer2> {
        let grid = parse_grid(input)?;
        let edges = Edge::ALL
            .into_iter()
            .flat_map(|d| {
                d.entry_points(&grid.clone())
                    .into_iter()
                    .map(move |p| (d, p))
            })
            .collect::<Vec<_>>();

        let mut scores = Vec::new();
        // let mut cache = HashMap::new();
        for beam in edges {
            let score = solve_grid(&grid, /* &mut cache, */ HashSet::new(), beam)
                .into_iter()
                .map(|(_, p)| p)
                .collect::<HashSet<_>>()
                .len();
            scores.push(score)
        }

        scores
            .into_iter()
            .max()
            .ok_or_else(|| anyhow::anyhow!("Failed to get any scores"))
    }
}

fn parse_grid(s: &str) -> anyhow::Result<Grid> {
    s.lines()
        .map(|v| v.chars().map(Node::try_from).try_collect::<Vec<_>>())
        .try_collect::<Grid>()
}

fn solve_grid(
    grid: &Grid,
    // cache: &mut HashMap<Beam, HashSet<Beam>>,
    mut visited: HashSet<Beam>,
    beam: Beam,
) -> HashSet<Beam> {
    if !(0..grid.len()).contains(&beam.1 .0)
        || !(0..grid[0].len()).contains(&beam.1 .1)
        || !visited.insert(beam)
    {
        visited
    // } else if let Some(n) = cache.get(&beam) {
    //     println!("{beam:?} (cache hit)");
    //     n.clone()
    } else {
        for v in handle_beam(grid, beam) {
            visited = solve_grid(grid, /* cache, */ visited, v);
        }

        // cache.insert(beam, visited.clone());

        visited
    }
}

fn handle_beam(grid: &Grid, (edge, (row, col)): Beam) -> Vec<Beam> {
    use Edge::*;
    use Node::*;

    let edges = match (edge, grid[row][col]) {
        (_, Empty)
        | (North, VertSplit)
        | (South, VertSplit)
        | (East, HorzSplit)
        | (West, HorzSplit) => vec![(edge)],
        (North, BackMirror) | (South, FwrdMirror) => vec![(West)],
        (North, FwrdMirror) | (South, BackMirror) => vec![(East)],
        (North, HorzSplit) | (South, HorzSplit) => vec![(East), (West)],
        (East, FwrdMirror) | (West, BackMirror) => vec![(North)],
        (East, BackMirror) | (West, FwrdMirror) => vec![(South)],
        (East, VertSplit) | (West, VertSplit) => vec![(North), (South)],
    };

    let mut beams = Vec::new();

    for edge in edges {
        if let Some(pos) = edge.with_position((row, col)) {
            beams.push((edge, pos));
        }
    }

    beams
}

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

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

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

    fn with_position(&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 entry_points(&self, grid: &Grid) -> Vec<Position> {
        match self {
            Edge::North => (0..grid[0].len()).map(|c| (grid.len() - 1, c)).collect(),
            Edge::East => (0..grid.len()).map(|r| (r, 0)).collect(),
            Edge::South => (0..grid[0].len()).map(|c| (0, c)).collect(),
            Edge::West => (0..grid.len()).map(|r| (r, grid[0].len() - 1)).collect(),
        }
    }
}

#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
enum Node {
    #[default]
    Empty,
    HorzSplit,
    VertSplit,
    FwrdMirror,
    BackMirror,
}

impl TryFrom<char> for Node {
    type Error = anyhow::Error;

    fn try_from(value: char) -> Result<Self, Self::Error> {
        Ok(match value {
            '.' => Node::Empty,
            '-' => Node::HorzSplit,
            '|' => Node::VertSplit,
            '/' => Node::FwrdMirror,
            '\\' => Node::BackMirror,
            c => anyhow::bail!("Invalid char: {c}"),
        })
    }
}

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

    const INPUT: &str = indoc::indoc! {r#"
        .|...\....
        |.-.\.....
        .....|-...
        ........|.
        ..........
        .........\
        ..../.\\..
        .-.-/..|..
        .|....-|.\
        ..//.|....
    "#};

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

    #[test]
    fn test_part_2() -> anyhow::Result<()> {
        Ok(assert_eq!(51, Day16::part_2(INPUT)?))
    }
}