use std::collections::{HashMap, HashSet};

#[derive(Copy, Clone, Debug, PartialEq, Eq)]
struct OrderRule {
    before: u8,
    after: u8,
}

impl OrderRule {
    fn follows_rule(self, update: &[u8]) -> bool {
        let mut before_index = None;
        let mut after_index = None;

        for (index, &page) in update.iter().enumerate() {
            if page == self.before {
                before_index = Some(index);
            } else if page == self.after {
                after_index = Some(index);
            }
        }

        if let (Some(bi), Some(ai)) = (before_index, after_index) {
            bi < ai
        } else {
            true
        }
    }
}

fn sort_by_order(order_rules: &[OrderRule], update: &[u8]) -> Vec<u8> {
    let pages: HashSet<u8> = update.iter().cloned().collect();
    let relevant_order_rules = order_rules
        .iter()
        .filter(|rule| pages.contains(&rule.before) && pages.contains(&rule.after))
        .collect::<Vec<_>>();

    //directed adjacency list
    let mut digraph = HashMap::new();
    for rule in relevant_order_rules {
        digraph
            .entry(rule.before)
            .or_insert(vec![rule.after])
            .push(rule.after)
    }

    topological_sort(digraph)
        .into_iter()
        .filter(|page| update.contains(page))
        .collect()
}

fn topological_sort(mut dag: HashMap<u8, Vec<u8>>) -> Vec<u8> {
    let mut sort = vec![];

    while !dag.is_empty() {
        let mut start_nodes: HashSet<u8> = dag.keys().cloned().collect();

        for destination in dag.values().flatten() {
            start_nodes.remove(destination);
        }
        assert_eq!(start_nodes.len(), 1);

        let start = start_nodes.into_iter().next().unwrap();
        sort.push(start);
        dag.remove(&start);
    }

    sort
}

pub fn day5_part1(input: &str) -> String {
    let (rules, updates) = parse(input);
    let sum: u32 = updates
        .iter()
        .filter(|update| rules.iter().all(|rule| rule.follows_rule(update)))
        .map(|update| update[update.len() / 2] as u32)
        .sum();
    sum.to_string()
}

pub fn day5_part2(input: &str) -> String {
    let (rules, updates) = parse(input);
    let sum: u32 = updates
        .iter()
        .filter(|update| rules.iter().any(|rule| !rule.follows_rule(update)))
        .map(|update| sort_by_order(&rules, update))
        .map(|update| update[update.len() / 2] as u32)
        .sum();
    sum.to_string()
}

fn parse(input: &str) -> (Vec<OrderRule>, Vec<Vec<u8>>) {
    let (rules, updates) = input.split_once("\n\n").unwrap();
    let rules = rules
        .lines()
        .map(|line| line.split_once('|').unwrap())
        .map(|(left, right)| OrderRule {
            before: left.parse().unwrap(),
            after: right.parse().unwrap(),
        })
        .collect();

    let updates = updates
        .lines()
        .map(|line| line.split(',').map(|page| page.parse().unwrap()).collect())
        .collect();

    (rules, updates)
}