use nom::{
    branch::alt,
    bytes::complete::tag,
    character::complete::{char as nom_char, line_ending, not_line_ending},
    combinator::value,
    multi::{many1, separated_list1},
    sequence::{delimited, pair},
    IResult,
};
use std::{cmp, collections::HashSet, fs};

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let input = fs::read_to_string("inputs/day_23.txt")?;
    let maze = parse_maze(&input).unwrap().1;
    dbg!(find_shortest_path(&maze).cost);
    Ok(())
}

fn find_shortest_path(maze: &Maze) -> Move {
    let mut visited = HashSet::new();
    visited.insert(maze.clone());
    let mut frontier = vec![Move {
        next_state: maze.clone(),
        cost: 0,
    }];

    let mut best_so_far: Option<Move> = None;

    while let Some(current) = frontier.pop() {
        if let Some(best_so_far) = &best_so_far {
            if current.cost >= best_so_far.cost {
                return best_so_far.clone();
            }
        }

        let next_moves: Vec<Move> = current
            .next_state
            .valid_moves()
            .into_iter()
            .map(|next| Move {
                cost: next.cost + current.cost,
                ..next
            })
            .collect();
        for next in next_moves {
            if next.next_state.is_complete() {
                best_so_far = if let Some(best_so_far) = best_so_far {
                    if best_so_far.cost < next.cost {
                        Some(best_so_far)
                    } else {
                        Some(next.clone())
                    }
                } else {
                    Some(next.clone())
                };
            } else if !visited.contains(&next.next_state) {
                visited.insert(next.next_state.clone());
                frontier.push(next);
            }
        }
        frontier.sort_unstable_by(|a, b| b.cost.cmp(&a.cost));
    }
    best_so_far.expect("There is no path through!")
}

#[derive(Debug, Clone)]
struct Move {
    next_state: Maze,
    cost: usize,
}

#[derive(Debug, Clone, PartialEq, Eq, Hash)]
struct Maze {
    corridor: Vec<Option<usize>>,
    rooms: Vec<Room>,
}

impl Maze {
    fn is_complete(&self) -> bool {
        self.rooms.iter().all(|room| room.is_complete())
    }

    fn valid_moves(&self) -> Vec<Move> {
        let mut valid_moves = Vec::new();

        for i in 0..self.corridor.len() {
            if let Some(shrimp) = self.corridor[i] {
                let target_room = &self.rooms[shrimp / 2 - 1];
                if target_room.can_enter() {
                    let route_free = (cmp::min(shrimp, i)..=cmp::max(shrimp, i))
                        .all(|route_i| route_i == i || self.corridor[route_i].is_none());
                    if route_free {
                        let mut next_state = self.clone();
                        next_state.corridor[i] = None;
                        let next_room = &mut next_state.rooms[shrimp / 2 - 1];
                        let room_depth = next_room
                            .max_free()
                            .expect("no space in room, but we checked!");
                        next_room.contents[room_depth] = Some(shrimp);
                        let distance = room_depth + 1 + cmp::max(shrimp, i) - cmp::min(shrimp, i);
                        let cost = calculate_cost(shrimp, distance);
                        valid_moves.push(Move { next_state, cost });
                    }
                }
            }
        }

        for (room_i, room) in self
            .rooms
            .iter()
            .enumerate()
            .filter(|(_, room)| !room.can_enter())
        {
            if let Some((room_depth, shrimp)) = room
                .contents
                .iter()
                .enumerate()
                .filter_map(|(room_depth, maybe_shrimp)| {
                    maybe_shrimp.map(|shrimp| (room_depth, shrimp))
                })
                .next()
            {
                for corridor_i in 0..self.corridor.len() {
                    let in_entrance = self.rooms.iter().any(|room| room.entrance == corridor_i);
                    let route_free = (cmp::min(room.entrance, corridor_i)
                        ..=cmp::max(room.entrance, corridor_i))
                        .all(|route_i| self.corridor[route_i].is_none());
                    if !in_entrance && route_free {
                        let mut next_state = self.clone();
                        next_state.corridor[corridor_i] = Some(shrimp);
                        next_state.rooms[room_i].contents[room_depth] = None;

                        let distance = room_depth + 1 + cmp::max(room.entrance, corridor_i)
                            - cmp::min(room.entrance, corridor_i);
                        let cost = calculate_cost(shrimp, distance);
                        valid_moves.push(Move { next_state, cost });
                    }
                }
            }
        }

        valid_moves
    }
}

fn calculate_cost(shrimp: usize, distance: usize) -> usize {
    let shrimp_cost = match shrimp {
        2 => 1,
        4 => 10,
        6 => 100,
        8 => 1000,
        _ => panic!("Unknown shrimp"),
    };
    shrimp_cost * distance
}

#[derive(Debug, Clone, PartialEq, Eq, Hash)]
struct Room {
    entrance: usize,
    contents: Vec<Option<usize>>,
}

impl Room {
    fn is_complete(&self) -> bool {
        self.contents
            .iter()
            .all(|slot| slot == &Some(self.entrance))
    }
    fn can_enter(&self) -> bool {
        self.contents
            .iter()
            .all(|slot| slot.is_none() || slot == &Some(self.entrance))
    }
    fn max_free(&self) -> Option<usize> {
        self.contents.iter().rposition(|slot| slot.is_none())
    }
}

fn parse_maze(input: &str) -> IResult<&str, Maze> {
    let (input, _) = pair(not_line_ending, line_ending)(input)?; // skip first line
    let (input, corridor) = delimited(nom_char('#'), many1(corridor_contents), tag("#\n"))(input)?;

    let (input, rooms_line_1) = delimited(
        tag("###"),
        separated_list1(nom_char('#'), corridor_contents),
        tag("###\n"),
    )(input)?;
    let (input, rooms_line_2) = delimited(
        tag("  #"),
        separated_list1(nom_char('#'), corridor_contents),
        tag("#\n"),
    )(input)?;

    let rooms = vec![
        Room {
            entrance: 2,
            contents: vec![rooms_line_1[0], Some(8), Some(8), rooms_line_2[0]],
        },
        Room {
            entrance: 4,
            contents: vec![rooms_line_1[1], Some(6), Some(4), rooms_line_2[1]],
        },
        Room {
            entrance: 6,
            contents: vec![rooms_line_1[2], Some(4), Some(2), rooms_line_2[2]],
        },
        Room {
            entrance: 8,
            contents: vec![rooms_line_1[3], Some(2), Some(6), rooms_line_2[3]],
        },
    ];

    Ok((input, Maze { corridor, rooms }))
}

fn corridor_contents(input: &str) -> IResult<&str, Option<usize>> {
    alt((
        value(None, nom_char('.')),
        value(Some(2), nom_char('A')),
        value(Some(4), nom_char('B')),
        value(Some(6), nom_char('C')),
        value(Some(8), nom_char('D')),
    ))(input)
}