use nom::{
    branch::alt,
    bytes::complete::tag,
    character::complete::{alpha1, i64, line_ending},
    combinator::map,
    multi::separated_list1,
    sequence::tuple,
    IResult,
};
use std::{
    collections::{BTreeMap, BTreeSet},
    fs,
};

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let input = fs::read_to_string("inputs/day_21.txt")?;
    let mut monkeys = Monkeys::parser(&input).unwrap().1;

    dbg!(monkeys.eval("root"));

    let (original_root_left, original_root_right) =
        monkeys.0.get("root").expect("No root monkey!").values();

    monkeys.0.insert(
        "root".into(),
        Expression::Eql(original_root_left, original_root_right.expect("foo")),
    );
    monkeys
        .0
        .insert("humn".into(), Expression::Value(Value::Input));

    while monkeys.0.len() > 1 {
        monkeys.simplify();
    }
    println!("{}", monkeys.print_expression("root"));

    Ok(())
}

#[derive(Debug, Clone)]
struct Monkeys(BTreeMap<String, Expression>);

#[derive(Debug, Clone)]
enum Expression {
    Value(Value),
    Add(Value, Value),
    Sub(Value, Value),
    Mul(Value, Value),
    Div(Value, Value),
    Eql(Value, Value),
}

#[derive(Debug, Clone)]
enum Value {
    Ref(String),
    Literal(i64),
    Input,
}

impl Monkeys {
    fn parser(input: &str) -> IResult<&str, Self> {
        map(
            separated_list1(line_ending, tuple((alpha1, tag(": "), Expression::parser))),
            |lines| {
                Monkeys(
                    lines
                        .into_iter()
                        .map(|(id, _, expression): (&str, _, Expression)| {
                            (id.to_owned(), expression)
                        })
                        .collect(),
                )
            },
        )(input)
    }

    fn simplify(&mut self) {
        let ids: Vec<String> = self.0.keys().cloned().collect();

        // simplify to literals where possible
        for id in &ids {
            if let Some(val) = self.eval(id) {
                self.0
                    .insert(id.clone(), Expression::Value(Value::Literal(val)));
            }
        }

        // simplify single sides of expressions where possible (especially unwrap refs)
        for id in &ids {
            let monkey = self.0.get(id).expect("Unknown monkey!");
            let simplified = match monkey {
                Expression::Value(v) => Expression::Value(self.simplify_value(v)),
                Expression::Add(a, b) => {
                    Expression::Add(self.simplify_value(a), self.simplify_value(b))
                }
                Expression::Sub(a, b) => {
                    Expression::Sub(self.simplify_value(a), self.simplify_value(b))
                }
                Expression::Mul(a, b) => {
                    Expression::Mul(self.simplify_value(a), self.simplify_value(b))
                }
                Expression::Div(a, b) => {
                    Expression::Div(self.simplify_value(a), self.simplify_value(b))
                }
                Expression::Eql(a, b) => {
                    Expression::Eql(self.simplify_value(a), self.simplify_value(b))
                }
            };
            self.0.insert(id.clone(), simplified);

            if let Some(val) = self.eval(id) {
                self.0
                    .insert(id.clone(), Expression::Value(Value::Literal(val)));
            }
        }

        // simplify across the equals
        for id in &ids {
            let monkey = self.0.get(id).expect("Unknown monkey!").clone();
            match monkey {
                Expression::Eql(Value::Ref(r), Value::Literal(v))
                | Expression::Eql(Value::Literal(v), Value::Ref(r)) => {
                    let referenced_monkey = self.0.get(&r).expect("Unknown monkey!").clone();

                    let new_eql = match referenced_monkey {
                        Expression::Add(new_lhs, Value::Literal(added))
                        | Expression::Add(Value::Literal(added), new_lhs) => {
                            Some((new_lhs.clone(), v - added))
                        }
                        Expression::Sub(new_lhs, Value::Literal(subtracted)) => {
                            Some((new_lhs.clone(), v + subtracted))
                        }
                        Expression::Sub(Value::Literal(subtracted_from), new_lhs) => {
                            Some((new_lhs.clone(), subtracted_from - v))
                        }
                        Expression::Mul(new_lhs, Value::Literal(mult))
                        | Expression::Mul(Value::Literal(mult), new_lhs) => {
                            Some((new_lhs.clone(), v / mult))
                        }
                        Expression::Div(new_lhs, Value::Literal(denom)) => {
                            Some((new_lhs.clone(), v * denom))
                        }
                        _ => None,
                    };

                    if let Some((new_lhs, new_v)) = new_eql {
                        self.0
                            .insert(id.clone(), Expression::Eql(new_lhs, Value::Literal(new_v)));
                    }
                }
                _ => {}
            }
        }

        // clearing out unreferenced monkeys
        let monkey_references: BTreeSet<String> = self
            .0
            .values()
            .flat_map(|expression| match expression.values() {
                (Value::Ref(v1), Some(Value::Ref(v2))) => vec![v1, v2],
                (Value::Ref(v), _) | (_, Some(Value::Ref(v))) => vec![v],
                _ => vec![],
            })
            .collect();
        for id in &ids {
            if id != "root" && !monkey_references.contains(id) {
                self.0.remove(id);
            }
        }
    }

    fn simplify_value(&self, val: &Value) -> Value {
        match val {
            Value::Ref(r) => {
                let referenced = self.0.get(r).expect("Unknown monkey!");
                match referenced {
                    Expression::Value(v) => v.clone(),
                    _ => val.clone(),
                }
            }
            v => v.clone(),
        }
    }

    fn eval(&self, id: &str) -> Option<i64> {
        let monkey = self.0.get(id).expect("Unknown monkey!");

        match monkey {
            Expression::Value(v) => self.resolve_value(v),
            Expression::Add(a, b) => self
                .resolve_value(a)
                .zip(self.resolve_value(b))
                .map(|(a, b)| a + b),
            Expression::Sub(a, b) => self
                .resolve_value(a)
                .zip(self.resolve_value(b))
                .map(|(a, b)| a - b),
            Expression::Mul(a, b) => self
                .resolve_value(a)
                .zip(self.resolve_value(b))
                .map(|(a, b)| a * b),
            Expression::Div(a, b) => self
                .resolve_value(a)
                .zip(self.resolve_value(b))
                .map(|(a, b)| a / b),
            Expression::Eql(a, b) => self
                .resolve_value(a)
                .zip(self.resolve_value(b))
                .map(|(a, b)| if a == b { 1 } else { 0 }),
        }
    }

    fn resolve_value(&self, value: &Value) -> Option<i64> {
        match value {
            Value::Ref(id) => self.eval(id),
            Value::Literal(num) => Some(*num),
            Value::Input => None,
        }
    }

    fn print_expression(&self, id: &str) -> String {
        let monkey = self.0.get(id).expect("Unknown monkey!");

        // The crazy newlines here make it easier to see the
        // brackets. Drop it in emacs and autoindent, and it's easy to
        // see the next expression which isn't being collapsed.
        match monkey {
            Expression::Value(v) => format!("(\n{}\n)", self.print_value(v)),
            Expression::Add(a, b) => {
                format!("(\n{}\n+\n{}\n)", self.print_value(a), self.print_value(b))
            }
            Expression::Sub(a, b) => {
                format!("(\n{}\n-\n{}\n)", self.print_value(a), self.print_value(b))
            }
            Expression::Mul(a, b) => {
                format!("(\n{}\n*\n{}\n)", self.print_value(a), self.print_value(b))
            }
            Expression::Div(a, b) => {
                format!("(\n{}\n/\n{}\n)", self.print_value(a), self.print_value(b))
            }
            Expression::Eql(a, b) => {
                format!("(\n{}\n=\n{}\n)", self.print_value(a), self.print_value(b))
            }
        }
    }

    fn print_value(&self, value: &Value) -> String {
        match value {
            Value::Ref(id) => self.print_expression(id),
            Value::Literal(num) => format!("{}", num),
            Value::Input => "input".into(),
        }
    }
}

impl Expression {
    fn parser(input: &str) -> IResult<&str, Self> {
        alt((
            map(
                tuple((Value::parser, tag(" + "), Value::parser)),
                |(a, _, b)| Expression::Add(a, b),
            ),
            map(
                tuple((Value::parser, tag(" - "), Value::parser)),
                |(a, _, b)| Expression::Sub(a, b),
            ),
            map(
                tuple((Value::parser, tag(" * "), Value::parser)),
                |(a, _, b)| Expression::Mul(a, b),
            ),
            map(
                tuple((Value::parser, tag(" / "), Value::parser)),
                |(a, _, b)| Expression::Div(a, b),
            ),
            map(Value::parser, Expression::Value),
        ))(input)
    }

    fn values(&self) -> (Value, Option<Value>) {
        match self.clone() {
            Expression::Value(v) => (v, None),
            Expression::Add(a, b) => (a, Some(b)),
            Expression::Sub(a, b) => (a, Some(b)),
            Expression::Mul(a, b) => (a, Some(b)),
            Expression::Div(a, b) => (a, Some(b)),
            Expression::Eql(a, b) => (a, Some(b)),
        }
    }
}

impl Value {
    fn parser(input: &str) -> IResult<&str, Self> {
        alt((
            map(i64, Value::Literal),
            map(alpha1, |s: &str| Value::Ref(s.to_owned())),
        ))(input)
    }
}