use nom::{
    branch::alt,
    bytes::complete::tag,
    character::complete::{alpha1, line_ending, u32},
    combinator::map,
    multi::separated_list1,
    sequence::tuple,
    IResult,
};
use std::{
    collections::{BTreeMap, BTreeSet},
    fs,
};

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let input = fs::read_to_string("inputs/day_16.txt")?;
    let nodes = Nodes::parser(&input).unwrap().1;
    let mut condensed = nodes.condense();

    let mut initial_open_valves = BTreeSet::new();
    initial_open_valves.insert(0);

    let initial_state = State {
        actors: vec![Actor {
            position: 0,
            time_remaining: 30,
        }],
        open_valves: initial_open_valves.clone(),
    };

    dbg!(condensed.find_optimal_pressure_relieved(&initial_state));

    let initial_state_with_elephant = State {
        actors: vec![
            Actor {
                position: 0,
                time_remaining: 26
            };
            2
        ],
        open_valves: initial_open_valves,
    };
    dbg!(condensed.find_optimal_pressure_relieved(&initial_state_with_elephant));

    Ok(())
}

#[derive(Debug, Clone)]
struct Nodes {
    nodes: BTreeMap<String, Node>,
}

#[derive(Debug, Clone)]
struct Node {
    id: String,
    flow_rate: u32,
    exits: BTreeMap<String, u32>,
}

#[derive(Debug)]
struct CondensedNodes {
    nodes: Vec<CondensedNode>,
    cache: BTreeMap<State, u32>,
}

#[derive(Debug, Clone)]
struct CondensedNode {
    flow_rate: u32,
    exits: Vec<u32>,
}

impl Nodes {
    fn parser(input: &str) -> IResult<&str, Self> {
        map(separated_list1(line_ending, Node::parser), |nodes| Nodes {
            nodes: nodes
                .into_iter()
                .map(|node| (node.id.clone(), node))
                .collect(),
        })(input)
    }

    fn condense(&self) -> CondensedNodes {
        let node_ids: Vec<String> = self
            .nodes
            .values()
            .filter(|n| n.id == "AA" || n.flow_rate > 0)
            .map(|n| n.id.clone())
            .collect();

        let mut condensed = CondensedNodes {
            nodes: self
                .nodes
                .values()
                .filter(|n| n.id == "AA" || n.flow_rate > 0)
                .map(|n| CondensedNode {
                    flow_rate: n.flow_rate,
                    exits: Vec::new(),
                })
                .collect(),
            cache: BTreeMap::new(),
        };

        for (id, mut node) in condensed.nodes.iter_mut().enumerate() {
            node.exits = node_ids
                .iter()
                .map(|original_destination_id| {
                    // +1 because in condensed it includes opening the valve
                    self.find_shortest_path(&node_ids[id], &original_destination_id) + 1
                })
                .collect()
        }

        condensed
    }

    fn find_shortest_path(&self, from: &str, to: &str) -> u32 {
        let mut frontier: BTreeSet<&str> = BTreeSet::new();
        let mut explored: BTreeSet<&str> = BTreeSet::new();
        let mut distance = 0;

        explored.insert(from);
        frontier.insert(from);

        while !frontier.is_empty() {
            let mut next_frontier: BTreeSet<&str> = BTreeSet::new();
            distance += 1;

            for frontier_point in frontier {
                for adjacent_point in self.nodes.get(frontier_point).unwrap().exits.keys() {
                    if adjacent_point == to {
                        return distance;
                    }
                    if !explored.contains(&adjacent_point.as_ref()) {
                        explored.insert(adjacent_point);
                        next_frontier.insert(adjacent_point);
                    }
                }
            }

            frontier = next_frontier;
        }
        panic!("Didn't reach end");
    }
}

impl Node {
    fn parser(input: &str) -> IResult<&str, Self> {
        map(
            tuple((
                tag("Valve "),
                alpha1,
                tag(" has flow rate="),
                u32,
                alt((
                    tag("; tunnels lead to valves "),
                    tag("; tunnel leads to valve "),
                )),
                separated_list1(tag(", "), alpha1),
            )),
            |(_, id, _, flow_rate, _, exits): (_, &str, _, u32, _, Vec<&str>)| Node {
                id: id.to_owned(),
                flow_rate,
                exits: exits
                    .into_iter()
                    .map(|destination| (destination.to_owned(), 1))
                    .collect(),
            },
        )(input)
    }
}

#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
struct State {
    actors: Vec<Actor>,
    open_valves: BTreeSet<usize>,
}

#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
struct Actor {
    position: usize,
    time_remaining: u32,
}

impl State {
    fn sort_actors(&mut self) {
        self.actors.sort()
    }
}

impl CondensedNodes {
    fn find_optimal_pressure_relieved(&mut self, state: &State) -> u32 {
        let cache_value = self.cache.get(state);
        if let Some(cache_value) = cache_value {
            return *cache_value;
        }

        let mut computed_value = 0;
        for (actor_i, actor) in state.actors.iter().enumerate() {
            let exits_to_investigate: Vec<(usize, u32)> = self.nodes[actor.position]
                .exits
                .iter()
                .enumerate()
                .filter(|(destination, distance)| {
                    !state.open_valves.contains(destination) && **distance <= actor.time_remaining
                })
                .map(|(destination, distance)| (destination.clone(), distance.clone()))
                .collect();

            for (destination, distance) in exits_to_investigate {
                let mut open_valves = state.open_valves.clone();
                open_valves.insert(destination.clone());
                let new_actor = Actor {
                    time_remaining: actor.time_remaining - distance,
                    position: destination.clone(),
                };
                let relief_from_this_valve =
                    self.nodes[destination].flow_rate * new_actor.time_remaining as u32;
                let mut actors = state.actors.clone();
                actors[actor_i] = new_actor;

                let mut state_this_way = State {
                    actors,
                    open_valves,
                };
                state_this_way.sort_actors();

                let recursive_relief = self.find_optimal_pressure_relieved(&state_this_way);
                let relief = relief_from_this_valve + recursive_relief;
                computed_value = computed_value.max(relief);
            }
        }

        self.cache.insert(state.clone(), computed_value);
        computed_value
    }
}