use aoc_runner_derive::{aoc, aoc_generator};
use std::collections::{HashMap, HashSet};

type Edges = HashMap<String, Vec<String>>;

#[aoc_generator(day11)]
fn parse(input: &str) -> Edges {
    let mut edges: Edges = HashMap::new();
    for l in input.lines() {
        let (k, rest) = l.split_once(": ").unwrap();
        for v in rest.split_ascii_whitespace() {
            edges.entry(k.to_string()).or_default().push(v.to_string());
        }
    }

    edges
}

fn find_path(cur: &str, goal: &str, edges: &Edges) -> u64 {
    if cur == goal {
        return 1;
    }
    if let Some(nexts) = edges.get(cur) {
        nexts.iter().map(|n| find_path(n, goal, edges)).sum()
    } else {
        0
    }
}

fn mark_paths<'a>(
    cur: &'a str,
    edges: &'a Edges,
    mut reachable: HashSet<&'a str>,
) -> HashSet<&'a str> {
    reachable.insert(cur);
    if let Some(nexts) = edges.get(cur) {
        for n in nexts {
            if !reachable.contains(&n.as_str()) {
                reachable = mark_paths(n, edges, reachable);
            }
        }
    }
    reachable
}

#[aoc(day11, part1)]
fn part1(fwd: &Edges) -> u64 {
    find_path("you", "out", fwd)
}

#[aoc(day11, part2)]
fn part2(edges: &Edges) -> u64 {
    let mut rev = Edges::from_iter(edges.keys().map(|k| (k.to_owned(), vec![])));
    for (from, tos) in edges {
        for to in tos {
            rev.entry(to.to_owned()).or_default().push(from.to_owned());
        }
    }
    let mut reachable_dac = mark_paths("dac", edges, HashSet::new());
    reachable_dac = mark_paths("dac", &rev, reachable_dac);
    let mut reachable_fft = mark_paths("fft", edges, HashSet::new());
    reachable_fft = mark_paths("fft", &rev, reachable_fft);

    let reachable: HashSet<&str> = reachable_dac
        .intersection(&reachable_fft)
        .copied()
        .collect();

    let unreachable: HashSet<&str> = HashSet::from_iter(edges.keys().map(|k| k.as_str()))
        .difference(&reachable)
        .copied()
        .collect();

    let mut reachable_edges = edges.clone();

    for k in &unreachable {
        reachable_edges.remove(*k);
    }
    for (_k, v) in reachable_edges.iter_mut() {
        for ur in &unreachable {
            if let Some(idx) = v.iter().position(|s| s == ur) {
                v.remove(idx);
            }
        }
    }

    // This is a bit of a cheat from viewing the graph and realizing all paths are ordered this way.
    // However, knowing that fact that we can partition in this way, we could use a shortest-path algo
    // to determine which order fft and dac are in.
    //
    // The assumption holds on my data and the example, so I think it's probably true of all inputs.
    find_path("svr", "fft", &reachable_edges)
        * find_path("fft", "dac", &reachable_edges)
        * find_path("dac", "out", &reachable_edges)
}

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

    use super::*;

    const EXAMPLE: &str = "aaa: you hhh
you: bbb ccc
bbb: ddd eee
ccc: ddd eee fff
ddd: ggg
eee: out
fff: out
ggg: out
hhh: ccc fff iii
iii: out";

    const EXAMPLE2: &str = "svr: aaa bbb
aaa: fft
fft: ccc
bbb: tty
tty: ccc
ccc: ddd eee
ddd: hub
hub: fff
eee: dac
dac: fff
fff: ggg hhh
ggg: out
hhh: out";

    #[rstest]
    #[case(EXAMPLE, 5)]
    fn part1_example(#[case] input: &str, #[case] expected: u64) {
        assert_eq!(part1(&parse(input)), expected);
    }

    #[rstest]
    #[case(EXAMPLE2, 2)]
    fn part2_example(#[case] input: &str, #[case] expected: u64) {
        assert_eq!(part2(&parse(input)), expected);
    }
}