use nom::{
    branch::alt,
    bytes::complete::tag,
    character::complete::{alpha1, char as nom_char, line_ending},
    combinator::{map, value},
    multi::{many1, separated_list1},
    sequence::{pair, preceded, separated_pair, terminated},
    IResult,
};
use std::{collections::BTreeMap, fs, ops::Range};

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let input = fs::read_to_string("inputs/day_8.txt")?;
    let mut directions = Directions::parser(&input).unwrap().1;
    dbg!(directions.steps_from_a_to_z());
    dbg!(directions.ghost_steps_from_a_to_z());

    Ok(())
}

#[derive(Debug)]
struct Directions {
    turns: Vec<Turn>,
    packed_map: Vec<PackedFork>,
    packed_ghost_starts: Range<u16>,
    packed_ghost_destinations: Range<u16>,
    distance_to_ghost_dest_cache: BTreeMap<(u16, usize), (u16, usize)>,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum Turn {
    Left,
    Right,
}

#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone)]
struct Location(String);

#[derive(Debug)]
struct Fork {
    left: Location,
    right: Location,
}

#[derive(Debug)]
struct PackedFork {
    left: u16,
    right: u16,
}

impl Directions {
    fn parser(input: &str) -> IResult<&str, Self> {
        map(
            separated_pair(
                many1(Turn::parser),
                pair(line_ending, line_ending),
                separated_list1(
                    line_ending,
                    separated_pair(Location::parser, tag(" = "), Fork::parser),
                ),
            ),
            |(turns, map)| {
                let map: BTreeMap<Location, Fork> = map.into_iter().collect();
                let mut locations: Vec<Location> = map.keys().cloned().collect();
                locations.sort_by_key(|l| l.0.chars().rev().collect::<String>());

                Directions {
                    turns,
                    packed_map: locations
                        .iter()
                        .map(|l| {
                            let unpacked_fork = map.get(l).unwrap();
                            PackedFork {
                                left: locations
                                    .iter()
                                    .position(|f| *f == unpacked_fork.left)
                                    .unwrap() as u16,
                                right: locations
                                    .iter()
                                    .position(|f| *f == unpacked_fork.right)
                                    .unwrap() as u16,
                            }
                        })
                        .collect(),
                    packed_ghost_starts: 0
                        ..locations.iter().position(|l| !l.ghost_start()).unwrap() as u16,
                    packed_ghost_destinations: locations.iter().position(|l| l.ghost_end()).unwrap()
                        as u16
                        ..locations.len() as u16,
                    distance_to_ghost_dest_cache: BTreeMap::new(),
                }
            },
        )(input)
    }

    fn steps_from_a_to_z(&self) -> usize {
        let mut step_count = 0;
        let mut current_location: u16 = 0;

        for dir in self.turns.iter().cycle() {
            let current_fork: &PackedFork = &self.packed_map[current_location as usize];
            current_location = match dir {
                Turn::Left => current_fork.left,
                Turn::Right => current_fork.right,
            };
            step_count += 1;

            if current_location == self.packed_map.len() as u16 - 1 {
                return step_count;
            }
        }
        unreachable!()
    }

    fn ghost_steps_from_a_to_z(&mut self) -> usize {
        let mut current_locations: Vec<(u16, usize)> = self
            .packed_ghost_starts
            .clone()
            .map(|start| self.ghost_step_to_next_end(start, 0))
            .collect();

        let mut any_stepped = true;
        while any_stepped {
            any_stepped = false;

            let current_max = current_locations
                .iter()
                .map(|(_, steps)| steps.clone())
                .max()
                .unwrap();

            for current_location in &mut current_locations {
                if current_location.1 < current_max {
                    any_stepped = true;
                    let (new_location, extra_steps) =
                        self.ghost_step_to_next_end(current_location.0, current_location.1);
                    current_location.0 = new_location;
                    current_location.1 += extra_steps;
                }
            }
        }

        current_locations[0].1
    }

    fn ghost_step_to_next_end(&mut self, start: u16, current_step_count: usize) -> (u16, usize) {
        self.distance_to_ghost_dest_cache
            .entry((start, current_step_count % self.turns.len()))
            .or_insert_with(|| {
                let mut step_count = 0;
                let mut current_location: u16 = start;

                for dir in self
                    .turns
                    .iter()
                    .skip(current_step_count % self.turns.len())
                    .cycle()
                {
                    let current_fork: &PackedFork = &self.packed_map[current_location as usize];
                    current_location = match dir {
                        Turn::Left => current_fork.left,
                        Turn::Right => current_fork.right,
                    };
                    step_count += 1;

                    if self.packed_ghost_destinations.contains(&current_location) {
                        break;
                    }
                }

                (current_location, step_count)
            })
            .clone()
    }
}

impl Turn {
    fn parser(input: &str) -> IResult<&str, Self> {
        alt((
            value(Turn::Left, nom_char('L')),
            value(Turn::Right, nom_char('R')),
        ))(input)
    }
}

impl Location {
    fn parser(input: &str) -> IResult<&str, Self> {
        map(alpha1, |l: &str| Location(l.to_string()))(input)
    }

    fn ghost_start(&self) -> bool {
        self.0.ends_with("A")
    }

    fn ghost_end(&self) -> bool {
        self.0.ends_with("Z")
    }
}

impl Fork {
    fn parser(input: &str) -> IResult<&str, Self> {
        map(
            preceded(
                tag("("),
                terminated(
                    separated_pair(Location::parser, tag(", "), Location::parser),
                    tag(")"),
                ),
            ),
            |(left, right)| Fork { left, right },
        )(input)
    }
}