use std::collections::HashMap;
use regex::{Regex, Match};

enum Rule {
    ACCEPTED,
    REJECTED,
    GOTO(String),
    GT(String, usize, String),
    LT(String, usize, String),
}

use aoc_runner_derive::{aoc, aoc_generator};
#[aoc_generator(day19, part1)]
fn parse(input: &str) -> (HashMap<String, Vec<Rule>>, Vec<HashMap<String, usize>>) {
    let (workflows_input, ratings_input) = input.split_once("\n\n").unwrap();
    let workflows = parse_workflows(workflows_input);

    // screw it im using regex
    let rer = Regex::new(r"^\{x=(\d+),m=(\d+),a=(\d+),s=(\d+)\}$").unwrap();
    let ratings: Vec<HashMap<String, usize>> = ratings_input.lines().map(|line| {
        let caps = rer.captures(line).unwrap();
        let x = parse_usize(caps.get(1));
        let m = parse_usize(caps.get(2));
        let a = parse_usize(caps.get(3));
        let s = parse_usize(caps.get(4));

        let mut vals = HashMap::new();
        vals.insert("x".to_string(), x);
        vals.insert("m".to_string(), m);
        vals.insert("a".to_string(), a);
        vals.insert("s".to_string(), s);
        vals
    }).collect();

    (workflows, ratings)
}

fn parse_usize(g: Option<Match>) -> usize {
    g.map_or(0, |m| m.as_str().parse().unwrap())
}

fn parse_workflows(workflows_input: &str) -> HashMap<String, Vec<Rule>> {
    let mut workflows: HashMap<String, Vec<Rule>> = HashMap::new();
    let rew = Regex::new(r"^(.+)\{(.+)\}$").unwrap();
    workflows_input.lines().for_each(|line| {
        let caps = rew.captures(line).unwrap();
        let key = caps.get(1).unwrap().as_str();
        let rules_str = caps.get(2).unwrap().as_str();

        let rules = rules_str.split(",").map(|expr| {
            if expr == "A" {
                return Rule::ACCEPTED;
            }
            if expr == "R" {
                return Rule::REJECTED;
            }
            if !expr.contains(":") {
                return Rule::GOTO(expr.to_string());
            }

            let (rule, to) = expr.split_once(":").unwrap();

            return if rule.contains(">") {
                let (prop, val) = rule.split_once(">").unwrap();
                Rule::GT(prop.to_string(), val.parse().unwrap(), to.to_string())
            } else if rule.contains("<") {
                let (prop, val) = rule.split_once("<").unwrap();
                Rule::LT(prop.to_string(), val.parse().unwrap(), to.to_string())
            } else {
                panic!("Unknown rule {}", rule)
            };
        }).collect();

        workflows.insert(key.to_string(), rules);
    });
    workflows
}

#[aoc(day19, part1)]
fn part1((workflows, ratings): &(HashMap<String, Vec<Rule>>, Vec<HashMap<String, usize>>)) -> usize {
    let mut accepted: Vec<HashMap<String, usize>> = vec!();
    'outer: for rating in ratings {
        let mut wf_key = "in";
        'middle: loop {
            if wf_key == "A" {
                accepted.push(rating.clone());
                continue 'outer;
            } else if wf_key == "R" {
                continue 'outer;
            }

            let rules = workflows.get(wf_key).unwrap();
            for rule in rules {
                match rule {
                    Rule::ACCEPTED => {
                        accepted.push(rating.clone());
                        continue 'outer;
                    }
                    Rule::REJECTED => {
                        continue 'outer;
                    }
                    Rule::GOTO(new_wf_key) => {
                        wf_key = new_wf_key;
                        continue 'middle;
                    }
                    Rule::GT(prop, val, to) => {
                        if rating.get(prop.as_str()).unwrap() > val {
                            wf_key = to;
                            continue 'middle;
                        }
                    }
                    Rule::LT(prop, val, to) => {
                        if rating.get(prop.as_str()).unwrap() < val {
                            wf_key = to;
                            continue 'middle;
                        }
                    }
                }
            }
        }
    }

    accepted.iter().map(|rating| rating.values().sum::<usize>()).sum::<usize>()
}


#[aoc_generator(day19, part2)]
fn parse_part2(input: &str) -> HashMap<String, Vec<Rule>> {
    let (workflows_input, _) = input.split_once("\n\n").unwrap();
    parse_workflows(workflows_input)
}

#[aoc(day19, part2)]
fn part2(workflows: &HashMap<String, Vec<Rule>>) -> usize {
    let mut stack: Vec<((usize, usize), (usize, usize), (usize, usize), (usize, usize), &str, usize)> =
        vec! {((1, 4000), (1, 4000), (1, 4000), (1, 4000), "in", 0)};
    let mut accepted: Vec<((usize, usize), (usize, usize), (usize, usize), (usize, usize))> = vec!();
    while let Some(range) = stack.pop() {
        let (x, m, a, s, wf_key, rule_key) = range;
        if wf_key == "A" {
            accepted.push((x, m, a, s));
            continue;
        } else if wf_key == "R" {
            continue;
        }

        if x.0 > x.1 || m.0 > m.1 || a.0 > a.1 || s.0 > s.1 { continue }

        let rules = workflows.get(wf_key).unwrap();
        let rule = &rules[rule_key];
        match rule {
            Rule::ACCEPTED => {
                accepted.push((x, m, a, s));
                continue;
            }
            Rule::REJECTED => {
                continue;
            }
            Rule::GOTO(new_wf_key) => {
                stack.push((x, m, a, s, new_wf_key, 0));
                continue;
            }
            Rule::GT(prop, val, to) => {
                match prop.as_str() {
                    "x" => {
                        stack.push(((val + 1, x.1), m, a, s, to.as_str(), 0));
                        stack.push(((x.0, *val), m, a, s, wf_key, rule_key + 1));
                    }
                    "m" => {
                        stack.push((x, (val + 1, m.1), a, s, to.as_str(), 0));
                        stack.push((x, (m.0, *val), a, s, wf_key, rule_key + 1));
                    }
                    "a" => {
                        stack.push((x, m, (val + 1, a.1), s, to.as_str(), 0));
                        stack.push((x, m, (a.0, *val), s, wf_key, rule_key + 1));
                    }
                    "s" => {
                        stack.push((x, m, a, (val + 1, s.1), to.as_str(), 0));
                        stack.push((x, m, a, (s.0, *val), wf_key, rule_key + 1));
                    }
                    _ => { panic!("unknown prop {}", prop) }
                }
            }
            Rule::LT(prop, val, to) => {
                match prop.as_str() {
                    "x" => {
                        stack.push(((x.0, val - 1), m, a, s, to.as_str(), 0));
                        stack.push(((*val, x.1), m, a, s, wf_key, rule_key + 1));
                    }
                    "m" => {
                        stack.push((x, (m.0, val - 1), a, s, to.as_str(), 0));
                        stack.push((x, (*val, m.1), a, s, wf_key, rule_key + 1));
                    }
                    "a" => {
                        stack.push((x, m, (a.0, val - 1), s, to.as_str(), 0));
                        stack.push((x, m, (*val, a.1), s, wf_key, rule_key + 1));
                    }
                    "s" => {
                        stack.push((x, m, a, (s.0, val - 1), to.as_str(), 0));
                        stack.push((x, m, a, (*val, s.1), wf_key, rule_key + 1));
                    }
                    _ => { panic!("unknown prop {}", prop) }
                }
            }
        }
    }

    accepted.iter().map(|(x, m, a, s)| {
        (x.1 - x.0 + 1) * (m.1 - m.0 + 1) * (a.1 - a.0 + 1) * (s.1 - s.0 + 1)
    }).sum()
}


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

    const EX: &str = r"px{a<2006:qkq,m>2090:A,rfg}
pv{a>1716:R,A}
lnx{m>1548:A,A}
rfg{s<537:gd,x>2440:R,A}
qs{s>3448:A,lnx}
qkq{x<1416:A,crn}
crn{x>2662:A,R}
in{s<1351:px,qqz}
qqz{s>2770:qs,m<1801:hdj,R}
gd{a>3333:R,R}
hdj{m>838:A,pv}

{x=787,m=2655,a=1222,s=2876}
{x=1679,m=44,a=2067,s=496}
{x=2036,m=264,a=79,s=2244}
{x=2461,m=1339,a=466,s=291}
{x=2127,m=1623,a=2188,s=1013}";

    #[test]
    fn part1_example() {
        assert_eq!(part1(&parse(EX)), 19114);
    }

    #[test]
    fn part2_example() {
        assert_eq!(part2(&parse_part2(EX)), 167409079868000);
    }
}