use regex::Regex;
use std::collections::{HashMap, HashSet};
use std::str::FromStr;

type Page = u32;

#[derive(Debug, Clone)]
struct Rule {
    before: Page,
    after: Page,
}

#[derive(Debug)]
enum ParseError {
    ParseError,
}

impl FromStr for Rule {
    type Err = ParseError;
    fn from_str(s: &str) -> Result<Self, Self::Err> {
        let Some(cs) = Regex::new("^([0-9]+)\\|([0-9]+)$").unwrap().captures(s) else {
            return Err(ParseError::ParseError);
        };
        let (_full, [bf, af]) = cs.extract();
        Ok(Rule {
            before: bf.parse().unwrap(),
            after: af.parse().unwrap(),
        })
    }
}

#[derive(Debug)]
struct RuleStatus {
    previous: HashMap<Page, HashSet<Page>>,
    invalid: HashSet<Page>,
}

impl RuleStatus {
    fn new(rules: &Vec<Rule>) -> Self {
        let mut previous = HashMap::new();
        for &Rule { before, after } in rules.iter() {
            match previous.get_mut(&after) {
                None => {
                    let mut set = HashSet::new();
                    set.insert(before);
                    previous.insert(after, set);
                }
                Some(set) => {
                    set.insert(before);
                }
            }
        }
        RuleStatus {
            previous,
            invalid: HashSet::new(),
        }
    }

    fn step(&mut self, page: &Page) -> bool {
        let RuleStatus { previous, invalid } = self;
        if invalid.contains(page) {
            return false;
        }
        if let Some(prev) = previous.get(page) {
            for p in prev.iter() {
                invalid.insert(*p);
            }
        }
        true
    }
}

type Update = Vec<Page>;

fn update_from_str(s: &str) -> Update {
    Regex::new("([0-9]+)")
        .unwrap()
        .find_iter(s)
        .map(|s| s.as_str().parse().unwrap())
        .collect()
}

fn is_update_valid(update: &Update, rules: &Vec<Rule>) -> bool {
    let mut status = RuleStatus::new(&rules);
    for page in update.iter() {
        if !status.step(page) {
            return false;
        }
    }
    true
}

fn middle_page(update: &Update) -> Page {
    update[update.len() / 2]
}

#[derive(Debug)]
struct Input {
    rules: Vec<Rule>,
    updates: Vec<Update>,
}

impl FromStr for Input {
    type Err = ParseError;
    fn from_str(s: &str) -> Result<Self, Self::Err> {
        let mut lines = s.lines();

        let mut input = Input {
            rules: vec![],
            updates: vec![],
        };

        // Read rules
        while let Some(l) = lines.next() {
            if l == "" {
                break;
            }
            input.rules.push(Rule::from_str(l).unwrap());
        }

        // Read updates
        while let Some(l) = lines.next() {
            input.updates.push(update_from_str(l));
        }

        Ok(input)
    }
}

fn sum_valid_updates(Input { rules, updates }: &Input) -> u32 {
    updates
        .iter()
        .filter(|u| is_update_valid(u, &rules))
        .map(|u| middle_page(u))
        .sum()
}

fn p1(input: &str) -> String {
    let input = Input::from_str(input).unwrap();
    sum_valid_updates(&input).to_string()
}

// Remove useless rules for this particular update
fn filter_rules(rules: &Vec<Rule>, update: &Update) -> Vec<Rule> {
    let updated: HashSet<&Page> = update.iter().collect();
    rules
        .iter()
        .cloned()
        .filter(|r| updated.contains(&r.before))
        .collect()
}

fn sort_rules(rules: &Vec<Rule>) -> Vec<Rule> {
    let mut sorted: Vec<Rule> = vec![];
    let mut afters: HashMap<Page, u8> = HashMap::new();
    for a in rules.iter().map(|r| r.after) {
        *afters.entry(a).or_default() += 1;
    }

    // Allocate new "droppable" rules by wrapping the originals into Options
    let mut rules: Vec<Option<&Rule>> = rules.iter().map(|r| Some(r)).collect();
    while rules.iter().any(|x| x.is_some()) {
        for rule in rules.iter_mut() {
            let Some(r) = *rule else { continue };
            if let None = afters.get(&r.before) {
                sorted.push(r.clone());
                match afters.get_mut(&r.after) {
                    None => {
                        panic!("This should not happen");
                    }
                    Some(1) => {
                        afters.remove(&r.after);
                    }
                    Some(x) => {
                        *x -= 1;
                    }
                };
                *rule = None;
            }
        }
    }

    sorted
}

fn update_reorder(update: &Update, rules: &Vec<Rule>) -> Option<Update> {
    let filtered_rules = filter_rules(rules, update);
    let sorted_rules = sort_rules(&filtered_rules);
    let mut result: Update = vec![];
    let mut pages: HashSet<Page> = update.iter().cloned().collect();
    for Rule { before, .. } in sorted_rules.iter() {
        if pages.contains(before) {
            result.push(*before);
            pages.remove(before);
        }
    }
    // Add unrestricted pages at the end
    for &p in pages.iter() {
        result.push(p);
    }
    if is_update_valid(&result, &sorted_rules) {
        Some(result)
    } else {
        None
    }
}

fn sum_reordered_updates(Input { rules, updates }: &Input) -> u32 {
    updates
        .iter()
        .filter(|u| !is_update_valid(u, &rules)) // drop valid updates
        .filter_map(|u| update_reorder(u, &rules)) // reorder updates
        .map(|u| middle_page(&u))
        .sum()
}

fn p2(input: &str) -> String {
    let input = Input::from_str(input).unwrap();
    sum_reordered_updates(&input).to_string()
}

fn main() {
    aoc2024::run_day("5", p1, p2);
}