use nom::{
    branch::alt,
    bytes::complete::tag,
    character::complete::{alpha1, line_ending, u32 as nom_u32},
    combinator::{map, value},
    multi::separated_list1,
    sequence::{pair, preceded, terminated, tuple},
    IResult,
};
use std::{collections::BTreeMap, fs, ops::Range};

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let input = fs::read_to_string("inputs/day_19.txt")?;
    let parsed = PartSortingMess::parser(&input).unwrap().1;
    dbg!(&parsed.accepted_part_ratings());
    dbg!(&parsed.count_distinct_combinations());

    Ok(())
}

#[derive(Debug)]
struct PartSortingMess {
    workflows: BTreeMap<String, Workflow>,
    parts: Vec<Part>,
}

#[derive(Debug)]
struct Workflow {
    id: String,
    conditions: Vec<WorkflowStep>,
    if_none_match: WorkflowOutcome,
}

#[derive(Debug)]
struct WorkflowStep {
    field: PartField,
    condition: WorkflowCondition,
    result: WorkflowOutcome,
}

#[derive(Debug, Clone, PartialEq, Eq)]
enum PartField {
    X,
    M,
    A,
    S,
}

#[derive(Debug)]
enum WorkflowCondition {
    LessThan(u32),
    GreaterThan(u32),
}

#[derive(Debug, Clone, PartialEq, Eq)]
enum WorkflowOutcome {
    Accept,
    Reject,
    Defer(String),
}

#[derive(Debug)]
struct Part {
    x: u32,
    m: u32,
    a: u32,
    s: u32,
}

#[derive(Debug, Clone)]
struct PartRange {
    x: Range<u32>,
    m: Range<u32>,
    a: Range<u32>,
    s: Range<u32>,
}

impl PartSortingMess {
    fn parser(input: &str) -> IResult<&str, Self> {
        map(
            tuple((
                separated_list1(line_ending, Workflow::parser),
                line_ending,
                line_ending,
                separated_list1(line_ending, Part::parser),
            )),
            |(workflows, _, _, parts)| PartSortingMess {
                workflows: workflows
                    .into_iter()
                    .map(|workflow| (workflow.id.clone(), workflow))
                    .collect(),
                parts,
            },
        )(input)
    }

    fn accepted_part_ratings(&self) -> u32 {
        let mut rating_sum = 0;

        for part in &self.parts {
            let mut outcome = WorkflowOutcome::Defer("in".into());
            while let Some(workflow_id) = outcome.next_workflow_id() {
                let workflow = self.workflows.get(&workflow_id).unwrap();
                outcome = workflow.process(part);
            }

            if outcome == WorkflowOutcome::Accept {
                rating_sum += part.rating();
            }
        }

        rating_sum
    }

    fn count_distinct_combinations(&self) -> u64 {
        self.count_combinations_for_workflow(
            "in",
            &PartRange {
                x: 1..4001,
                m: 1..4001,
                a: 1..4001,
                s: 1..4001,
            },
        )

        // potential workflow optimizations to make the stack not so deep
        // - eliminate conditions in a workflow where everything has the same outcome
        // - inline workflows with a single outcome
        // - eliminate unreachable branches like a<10, a<5
    }

    fn count_combinations_for_workflow(&self, workflow_id: &str, part_range: &PartRange) -> u64 {
        if part_range.combinations() == 0 {
            0
        } else {
            let mut combinations = 0;

            let workflow = self.workflows.get(workflow_id).unwrap();
            let mut remaining_range = part_range.clone();
            for condition in &workflow.conditions {
                let matched_range: PartRange;
                (matched_range, remaining_range) =
                    remaining_range.partition(&condition.field, &condition.condition);
                combinations += match &condition.result {
                    WorkflowOutcome::Accept => matched_range.combinations(),
                    WorkflowOutcome::Reject => 0,
                    WorkflowOutcome::Defer(next_workflow_id) => {
                        self.count_combinations_for_workflow(&next_workflow_id, &matched_range)
                    }
                };
            }
            combinations += match &workflow.if_none_match {
                WorkflowOutcome::Accept => remaining_range.combinations(),
                WorkflowOutcome::Reject => 0,
                WorkflowOutcome::Defer(next_workflow_id) => {
                    self.count_combinations_for_workflow(&next_workflow_id, &remaining_range)
                }
            };

            combinations
        }
    }
}

impl Workflow {
    fn parser(input: &str) -> IResult<&str, Self> {
        map(
            pair(
                alpha1,
                preceded(
                    tag("{"),
                    terminated(
                        pair(
                            separated_list1(tag(","), WorkflowStep::parser),
                            preceded(tag(","), WorkflowOutcome::parser),
                        ),
                        tag("}"),
                    ),
                ),
            ),
            |(id, (conditions, if_none_match)): (&str, _)| Workflow {
                id: id.to_owned(),
                conditions,
                if_none_match,
            },
        )(input)
    }

    fn process(&self, part: &Part) -> WorkflowOutcome {
        for condition in &self.conditions {
            let val = part.get(&condition.field);
            if condition.condition.matches(val) {
                return condition.result.clone();
            }
        }
        self.if_none_match.clone()
    }
}

impl WorkflowStep {
    fn parser(input: &str) -> IResult<&str, Self> {
        map(
            tuple((
                PartField::parser,
                WorkflowCondition::parser,
                preceded(tag(":"), WorkflowOutcome::parser),
            )),
            |(field, condition, result)| WorkflowStep {
                field,
                condition,
                result,
            },
        )(input)
    }
}

impl PartField {
    fn parser(input: &str) -> IResult<&str, Self> {
        alt((
            value(PartField::X, tag("x")),
            value(PartField::M, tag("m")),
            value(PartField::A, tag("a")),
            value(PartField::S, tag("s")),
        ))(input)
    }
}

impl WorkflowCondition {
    fn parser(input: &str) -> IResult<&str, Self> {
        alt((
            map(preceded(tag("<"), nom_u32), WorkflowCondition::LessThan),
            map(preceded(tag(">"), nom_u32), WorkflowCondition::GreaterThan),
        ))(input)
    }

    fn matches(&self, val: u32) -> bool {
        match self {
            WorkflowCondition::LessThan(me) => val < *me,
            WorkflowCondition::GreaterThan(me) => val > *me,
        }
    }

    fn partition(&self, val: &Range<u32>) -> (Range<u32>, Range<u32>) {
        match self {
            WorkflowCondition::LessThan(me) => (val.start..*me, *me..val.end),
            WorkflowCondition::GreaterThan(me) => (*me + 1..val.end, val.start..*me + 1),
        }
    }
}

impl WorkflowOutcome {
    fn parser(input: &str) -> IResult<&str, Self> {
        alt((
            value(WorkflowOutcome::Accept, tag("A")),
            value(WorkflowOutcome::Reject, tag("R")),
            map(alpha1, |id: &str| WorkflowOutcome::Defer(id.to_owned())),
        ))(input)
    }

    fn next_workflow_id(&self) -> Option<String> {
        match self {
            WorkflowOutcome::Accept | WorkflowOutcome::Reject => None,
            WorkflowOutcome::Defer(id) => Some(id.clone()),
        }
    }
}

impl Part {
    fn parser(input: &str) -> IResult<&str, Self> {
        map(
            tuple((
                tag("{x="),
                nom_u32,
                tag(",m="),
                nom_u32,
                tag(",a="),
                nom_u32,
                tag(",s="),
                nom_u32,
                tag("}"),
            )),
            |(_, x, _, m, _, a, _, s, _)| Part { x, m, a, s },
        )(input)
    }

    fn rating(&self) -> u32 {
        self.x + self.m + self.a + self.s
    }

    fn get(&self, field: &PartField) -> u32 {
        match field {
            PartField::X => self.x,
            PartField::M => self.m,
            PartField::A => self.a,
            PartField::S => self.s,
        }
    }
}

impl PartRange {
    fn combinations(&self) -> u64 {
        if self.x.is_empty() || self.m.is_empty() || self.a.is_empty() || self.s.is_empty() {
            0
        } else {
            (self.x.end - self.x.start) as u64
                * (self.m.end - self.m.start) as u64
                * (self.a.end - self.a.start) as u64
                * (self.s.end - self.s.start) as u64
        }
    }

    fn partition(
        &self,
        field: &PartField,
        condition: &WorkflowCondition,
    ) -> (PartRange, PartRange) {
        let (matched_range, unmatched_range) = match field {
            PartField::X => condition.partition(&self.x),
            PartField::M => condition.partition(&self.m),
            PartField::A => condition.partition(&self.a),
            PartField::S => condition.partition(&self.s),
        };

        let mut matched_part_range = self.clone();
        let mut unmatched_part_range = self.clone();

        match field {
            PartField::X => {
                matched_part_range.x = matched_range;
                unmatched_part_range.x = unmatched_range;
            }
            PartField::M => {
                matched_part_range.m = matched_range;
                unmatched_part_range.m = unmatched_range;
            }
            PartField::A => {
                matched_part_range.a = matched_range;
                unmatched_part_range.a = unmatched_range;
            }
            PartField::S => {
                matched_part_range.s = matched_range;
                unmatched_part_range.s = unmatched_range;
            }
        };

        (matched_part_range, unmatched_part_range)
    }
}