use std::collections::HashMap;
use std::fs::File;
use std::io::{BufRead, BufReader, Lines};
use std::ops::Range;
use std::time::Instant;

// BOILERPLATE
type InputIter = Lines<BufReader<File>>;

fn get_input() -> InputIter {
    let f = File::open("input").unwrap();
    let br = BufReader::new(f);
    br.lines()
}

fn main() {
    let start = Instant::now();
    let ans1 = problem1(get_input());
    let duration = start.elapsed();
    println!("Problem 1 solution: {} [{}s]", ans1, duration.as_secs_f64());

    let start = Instant::now();
    let ans2 = problem2(get_input());
    let duration = start.elapsed();
    println!("Problem 2 solution: {} [{}s]", ans2, duration.as_secs_f64());
}

// DATA

const INPUT_RANGE: Range<u64> = 1..4001;

fn empty_counters() -> HashMap<char, Vec<Range<u64>>> {
    HashMap::from([
        ('x', vec![INPUT_RANGE]),
        ('m', vec![INPUT_RANGE]),
        ('a', vec![INPUT_RANGE]),
        ('s', vec![INPUT_RANGE]),
    ])
}

#[derive(Debug)]
struct RulePredicate {
    op: PredicateOperator,
    var: char,
}

#[derive(Debug, Clone)]
enum RuleAction {
    Terminate(bool),
    Jump(String),
}

impl From<&str> for RuleAction {
    fn from(s: &str) -> Self {
        match s {
            "A" => Self::Terminate(true),
            "R" => Self::Terminate(false),
            s => Self::Jump(s.into()),
        }
    }
}

#[derive(Debug)]
struct Rule {
    pred: RulePredicate,
    action: RuleAction,
}

#[derive(Debug)]
enum PredicateOperator {
    Always,
    Lt(u64),
    Gt(u64),
}

fn range_overlap<T: num::Num + Ord + Copy>(r1: &Range<T>, r2: &Range<T>) -> Option<Range<T>> {
    let new_start = std::cmp::max(r1.start, r2.start);
    let new_end = std::cmp::min(r1.end - T::one(), r2.end - T::one());

    if new_start <= std::cmp::min(r1.end - T::one(), r2.end - T::one()) && new_end >= std::cmp::max(r1.start, r2.start)
    {
        Some(new_start..new_end + T::one())
    } else {
        None
    }
}

fn range_exclude<T: num::Num + Ord + Copy>(keep: &Range<T>, exclude: &Range<T>) -> Vec<Range<T>> {
    let mut residual = Vec::new();
    if let Some(overlap) = range_overlap(keep, exclude) {
        if keep.start < overlap.start {
            residual.push(keep.start..overlap.start);
        }
        if keep.end > overlap.end {
            residual.push(overlap.end..keep.end);
        }
    } else {
        residual.push(keep.clone());
    }
    residual
}

impl From<&str> for PredicateOperator {
    fn from(s: &str) -> Self {
        let (op_s, val_s) = s.split_at(1);
        match op_s {
            "<" => PredicateOperator::Lt(val_s.parse().unwrap()),
            ">" => PredicateOperator::Gt(val_s.parse().unwrap()),
            s => panic!("unknown operator {}", s),
        }
    }
}

impl From<&str> for RulePredicate {
    fn from(s: &str) -> Self {
        let (var_s, pred_s) = s.split_at(1);
        Self {
            op: pred_s.into(),
            var: var_s.chars().next().unwrap(),
        }
    }
}

impl RulePredicate {
    fn check(&self, part: &Part) -> bool {
        match self.op {
            PredicateOperator::Always => true,
            PredicateOperator::Gt(val) => part.0[&self.var] > val,
            PredicateOperator::Lt(val) => part.0[&self.var] < val,
        }
    }
    fn matching_range(&self) -> Range<u64> {
        let res = match self.op {
            PredicateOperator::Always => INPUT_RANGE,
            PredicateOperator::Gt(val) => val + 1..INPUT_RANGE.end,
            PredicateOperator::Lt(val) => INPUT_RANGE.start..val,
        };
        println!("   matching range for predicate {:?}: {:?}", self, res);
        res
    }
}

impl From<&str> for Rule {
    fn from(s: &str) -> Self {
        if let Some((predicate_s, action_s)) = s.split_once(':') {
            Self {
                pred: predicate_s.into(),
                action: action_s.into(),
            }
        } else {
            Self {
                pred: RulePredicate {
                    op: PredicateOperator::Always,
                    var: '\0',
                },
                action: s.into(),
            }
        }
    }
}

fn count_states(ranges: HashMap<char, Vec<Range<u64>>>) -> u64 {
    ['x', 'm', 'a', 's'].iter().map(|c| ranges[c].iter().map(|r| r.end - r.start).sum::<u64>()).product()
}

impl Rule {
    // Returns (matching_ranges, unmatching_ranges for next rule)
    fn possible_ranges(
        &self,
        wfs: &Workflows,
        ranges: HashMap<char, Vec<Range<u64>>>,
    ) -> (u64, HashMap<char, Vec<Range<u64>>>) {
        return match &self.action {
            RuleAction::Terminate(true) => {
                if let PredicateOperator::Always = self.pred.op {
                    // Always predicate is terminating and returns empty ranges
                    (count_states(ranges), empty_counters())
                } else {
                    // other predicates will pop up the stack and return unmatched ranges
                    let (mut matching, mut unmatching) = (ranges.clone(), ranges.clone());
                    if let Some(relevant_ranges) = ranges.get(&(self.pred.var)){
                        println!("  relevant: {:?}", relevant_ranges);
                        matching.insert(
                            self.pred.var,
                            relevant_ranges
                                .iter()
                                .filter_map(|range| range_overlap(range, &self.pred.matching_range()))
                                .collect(),
                        );
                        unmatching.insert(
                            self.pred.var,
                            relevant_ranges.iter().flat_map(|range| range_exclude(range, &self.pred.matching_range())).collect());
                        println!("  matching: {:?}", matching);
                        (count_states(matching), unmatching)
                    } else {
                        // relevant_ranges is empty so this is a failed state with no possibilities for one of the values
                        // probably we should never get here
                        (0, empty_counters())
                    }
                }
            }
            RuleAction::Terminate(false) => {
                if let PredicateOperator::Always = self.pred.op {
                    // Always predicate is terminating, with false returns 0 count and empty ranges
                    (0, empty_counters())
                } else {
                    let (mut matching, mut unmatching) = (ranges.clone(), ranges.clone());
                    if let Some(relevant_ranges) = ranges.get(&(self.pred.var)){
                        matching.insert(
                            self.pred.var,
                            relevant_ranges
                                .iter()
                                .filter_map(|range| range_overlap(range, &self.pred.matching_range()))
                                .collect(),
                        );
                        unmatching.insert(
                            self.pred.var,
                            relevant_ranges.iter().flat_map(|range| range_exclude(range, &self.pred.matching_range())).collect());
                        (0, unmatching)
                    } else {
                        // relevant_ranges is empty so this is a failed state with no possibilities for one of the values
                        // probably we should never get here
                        (0, empty_counters())
                    }
                }
            }
            RuleAction::Jump(wf) => {
                if let PredicateOperator::Always = self.pred.op {
                    // always predicate before a jump will always jump, so has no unmatching ranges
                    (wfs.0[wf].possible_ranges(wfs, ranges), empty_counters())
                } else {
                    let (mut matching, mut unmatching) = (ranges.clone(), ranges.clone());
                    if let Some(relevant_ranges) = ranges.get(&(self.pred.var)){
                        matching.insert(
                            self.pred.var,
                            relevant_ranges
                                .iter()
                                .filter_map(|range| range_overlap(range, &self.pred.matching_range()))
                                .collect(),
                        );
                        unmatching.insert(
                            self.pred.var,
                            relevant_ranges.iter().flat_map(|range| range_exclude(range, &self.pred.matching_range())).collect());
                        (wfs.0[wf].possible_ranges(wfs, matching), unmatching)
                    } else {
                        // no relevant ranges = no possible continuations
                        (0, empty_counters())
                    }
                }
            }
        };
    }
}

#[derive(Debug)]
struct Workflow {
    name: String,
    rules: Vec<Rule>,
}

impl From<&str> for Workflow {
    fn from(s: &str) -> Self {
        let (name_s, rest_s) = s.split_once('{').unwrap();
        let rules = rest_s.split_once('}').unwrap().0.split(',').map(|r| r.into()).collect();
        Self {
            name: name_s.into(),
            rules,
        }
    }
}
impl Workflow {
    fn execute(&self, part: &Part) -> RuleAction {
        for r in &self.rules {
            if r.pred.check(part) {
                return r.action.clone();
            }
        }
        panic!("unhandled part {:?}", part);
    }
    fn possible_ranges(
        &self,
        wfs: &Workflows,
        mut ranges: HashMap<char, Vec<Range<u64>>>,
    ) -> u64 {
        let mut accum = 0u64;
        println!("Entering {} with ranges {:?}", self.name, ranges);
        for r in &self.rules {
            println!(" evaluating rule: {:?} with {:?}", r, ranges);
            let (count, next_ranges) = r.possible_ranges(wfs, ranges);
            ranges = next_ranges;
            println!(" result of {:?}: count<{}> remaining<{:?}>", r, count, ranges);
            accum += count
        }
        println!("Count of {}: {}", self.name, accum);
        accum
    }
}

impl Workflows {
    fn execute(&self, part: &Part) -> bool {
        let mut action = RuleAction::Jump("in".into());
        loop {
            match &action {
                RuleAction::Terminate(b) => return *b,
                RuleAction::Jump(j) => {
                    let next_workflow = &self.0[j];
                    action = next_workflow.execute(part)
                }
            }
        }
    }

    fn count_possible_states(&self) -> u64 {
        let ranges = HashMap::from([
            ('x', vec![INPUT_RANGE]),
            ('m', vec![INPUT_RANGE]),
            ('a', vec![INPUT_RANGE]),
            ('s', vec![INPUT_RANGE]),
        ]);
        let possible_ranges = self.0["in".into()].possible_ranges(self, ranges);
        println!("possible_ranges: {:?}", possible_ranges);

        possible_ranges
    }
}

#[derive(Debug)]
struct Workflows(HashMap<String, Workflow>);

#[derive(Debug)]
struct Part(HashMap<char, u64>);

impl From<&str> for Part {
    fn from(s: &str) -> Self {
        let (_, vars_s) = s.split_once('{').unwrap();
        let vars = vars_s.split_once('}').unwrap().0.split(',');
        let mut part = HashMap::new();
        for var in vars {
            let (name, val) = var.split_once('=').unwrap();
            part.insert(name.chars().next().unwrap(), val.parse().unwrap());
        }
        Self(part)
    }
}

type Parts = Vec<Part>;

// PROBLEM 1 solution

fn problem1<T: BufRead>(mut input: Lines<T>) -> u64 {
    let mut wfs = Workflows(HashMap::new());
    let mut parts: Parts = Vec::new();
    while let Some(Ok(line)) = input.next() {
        if line != "" {
            let wf: Workflow = line.as_str().into();
            wfs.0.insert(wf.name.clone(), wf);
        } else {
            break;
        }
    }
    while let Some(Ok(line)) = input.next() {
        parts.push(line.as_str().into());
    }

    parts
        .iter()
        .filter(|part| wfs.execute(part))
        .map(|part| part.0.values().sum::<u64>())
        .sum::<u64>()
}

// PROBLEM 2 solution
fn problem2<T: BufRead>(mut input: Lines<T>) -> u64 {
    let mut wfs = Workflows(HashMap::new());
    let mut parts: Parts = Vec::new();
    while let Some(Ok(line)) = input.next() {
        if line != "" {
            let wf: Workflow = line.as_str().into();
            wfs.0.insert(wf.name.clone(), wf);
        } else {
            break;
        }
    }
    while let Some(Ok(line)) = input.next() {
        parts.push(line.as_str().into());
    }

    wfs.count_possible_states()
}

#[cfg(test)]
mod tests {
    use crate::*;
    use std::io::Cursor;
    const EXAMPLE: &str = &"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 problem1_example() {
        let c = Cursor::new(EXAMPLE);
        assert_eq!(problem1(c.lines()), 19114);
    }

    #[test]
    fn problem2_example() {
        let c = Cursor::new(EXAMPLE);
        assert_eq!(problem2(c.lines()), 167409079868000);
    }
}