#![forbid(elided_lifetimes_in_paths, unsafe_code)]

use aoc23::read;
use chumsky::prelude::*;
use std::collections::{HashMap, HashSet};

enum Instruction {
	Left,
	Right
}

impl Instruction {
	fn parser() -> impl Parser<char, Self, Error = Simple<char>> {
		choice((
			just('L').map(|_| Self::Left),
			just('R').map(|_| Self::Right)
		))
	}
}

#[derive(Debug)]
struct Node {
	left: String,
	right: String
}

fn parse_node_name() -> impl Parser<char, String, Error = Simple<char>> {
	none_of([',', ')', ' '])
		.repeated()
		.at_least(1)
		.map(|chars| chars.into_iter().collect())
}

impl Node {
	fn parser() -> impl Parser<char, Self, Error = Simple<char>> {
		just("(")
			.ignore_then(parse_node_name())
			.then_ignore(just(", "))
			.then(parse_node_name())
			.then_ignore(just(")"))
			.map(|(left, right)| Self { left, right })
	}
}

#[derive(Debug)]
struct Network {
	nodes: HashMap<String, Node>
}

impl Network {
	fn parser() -> impl Parser<char, Self, Error = Simple<char>> {
		parse_node_name()
			.then_ignore(just(" = "))
			.then(Node::parser())
			.then_ignore(just("\n"))
			.repeated()
			.map(|nodes| Self {
				nodes: nodes.into_iter().collect()
			})
	}

	fn node(&self, node: &str) -> &Node {
		self.nodes
			.get(node)
			.unwrap_or_else(|| panic!("Failed to get node {node:?}"))
	}

	fn left(&self, node: &str) -> &str {
		&self.node(node).left
	}

	fn right(&self, node: &str) -> &str {
		&self.node(node).right
	}
}

fn parser() -> impl Parser<char, (Vec<Instruction>, Network), Error = Simple<char>> {
	Instruction::parser()
		.repeated()
		.at_least(1)
		.then_ignore(just("\n\n"))
		.then(Network::parser())
		.then_ignore(end())
}

fn main() -> anyhow::Result<()> {
	let (instructions, network) = read("inputs/day8.txt", parser())?;

	let mut curr_node = "AAA";
	let mut i = 0;
	let mut steps = 0;
	while curr_node != "ZZZ" {
		curr_node = match instructions[i] {
			Instruction::Left => network.left(curr_node),
			Instruction::Right => network.right(curr_node)
		};
		i = (i + 1) % instructions.len();
		steps += 1;
	}
	println!("{steps}");

	let mut curr_nodes: HashSet<&str> = network
		.nodes
		.keys()
		.map(|node| node.as_str())
		.filter(|node| node.ends_with('A'))
		.collect();
	dbg!(curr_nodes.len());
	let mut i = 0;
	let mut steps = 0;
	while curr_nodes.iter().any(|node| !node.ends_with('Z')) {
		curr_nodes = curr_nodes
			.into_iter()
			.map(|node| match instructions[i] {
				Instruction::Left => network.left(node),
				Instruction::Right => network.right(node)
			})
			.collect();
		i = (i + 1) % instructions.len();
		steps += 1;
		if steps % 10000 == 0 {
			dbg!(steps);
		}
	}
	println!("{steps}");

	Ok(())
}