use std::error::Error;
use std::fs::File;
use std::io::{self, BufRead};
use std::vec::Vec;

fn read_paren_expression(input: &str) -> (&str, &str) {
    let mut open_count = 0;
    let mut close_count = 0;

    for i in 0..input.len() {
        match input.chars().nth(i).unwrap() {
            '(' => open_count += 1,
            ')' => {
                close_count += 1;
                if open_count == close_count {
                    return (&input[1..i], &input[i + 1..]);
                }
            }
            _ => {}
        }
    }

    panic!(
        "Encountered unablanced parentheses {} vs. {}",
        open_count, close_count
    );
}

fn read_number(input: &str) -> (&str, &str) {
    for i in 0..input.len() {
        if !input.chars().nth(i).unwrap().is_ascii_digit() {
            return (&input[0..i], &input[i..]);
        }
    }
    return (input, "");
}

#[derive(PartialEq, Clone, Debug)]
enum Expression {
    Num(u64),
    OpAdd,
    OpMul,
    SubExpr(Vec<Expression>),
}

fn parse_expression(input: &str) -> Vec<Expression> {
    let mut expr = Vec::new();
    let mut rest = input;

    while rest != "" {
        let c = rest.chars().nth(0).unwrap();
        if c.is_ascii_digit() {
            let (substr, tail) = read_number(rest);
            expr.push(Expression::Num(substr.parse().unwrap()));
            rest = tail;
        } else if c == '(' {
            let (substr, tail) = read_paren_expression(rest);
            expr.push(Expression::SubExpr(parse_expression(substr)));
            rest = tail;
        } else if c == '+' {
            expr.push(Expression::OpAdd);
            rest = &rest[1..];
        } else if c == '*' {
            expr.push(Expression::OpMul);
            rest = &rest[1..];
        } else {
            rest = &rest[1..];
        }
    }

    expr
}

fn eval_part1(expr: &Vec<Expression>) -> u64 {
    let subexpr_eval = expr.iter().map(|e| match e {
        Expression::SubExpr(subexpr) => Expression::Num(eval_part1(subexpr)),
        _ => e.clone(),
    });

    let mut res = 0;
    let mut op = Expression::OpAdd;
    for e in subexpr_eval {
        match e {
            Expression::Num(n) => match op {
                Expression::OpAdd => res += n,
                Expression::OpMul => res *= n,
                _ => panic!("Malformed expression. Not an op."),
            },
            Expression::OpAdd => op = Expression::OpAdd,
            Expression::OpMul => op = Expression::OpMul,
            _ => panic!("Malformed expression."),
        }
    }

    res
}

fn eval_part2(expr: &Vec<Expression>) -> u64 {
    let subexpr_eval: Vec<Expression> = expr
        .iter()
        .map(|e| match e {
            Expression::SubExpr(subexpr) => Expression::Num(eval_part2(subexpr)),
            _ => e.clone(),
        })
        .collect();

    let mut mul_eval: Vec<Expression> = Vec::new();
    let mut i = 0;
    while i < subexpr_eval.len() {
        match subexpr_eval[i] {
            Expression::OpAdd => {
                if let Expression::Num(n1) = mul_eval.pop().unwrap() {
                    if let Expression::Num(n2) = subexpr_eval[i + 1] {
                        mul_eval.push(Expression::Num(n1 + n2));
                        i += 2;
                    } else {
                        panic!("Malformed expression: Expected Number")
                    }
                } else {
                    panic!("Malformed expression: Expected Number")
                }
            }
            _ => {
                mul_eval.push(subexpr_eval[i].clone());
                i += 1;
            }
        }
    }

    let mut res = 1;

    for e in mul_eval {
        match e {
            Expression::Num(n) => res *= n,
            Expression::OpMul => {}
            _ => panic!("Malformed expression."),
        }
    }

    res
}

fn main() -> Result<(), Box<dyn Error>> {
    let file = File::open("inputs/day18.txt")?;
    let lines = io::BufReader::new(file).lines().map(|l| l.unwrap());

    let exprs: Vec<Vec<Expression>> = lines.map(|l| parse_expression(&l)).collect();

    let result1 = exprs.iter().map(|e| eval_part1(e)).fold(0, |x, y| x + y);

    println!("Answer Part 1: {}", result1);

    let result2 = exprs.iter().map(|e| eval_part2(e)).fold(0, |x, y| x + y);

    println!("Answer Part 2: {}", result2);

    Ok(())
}

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

    #[test]
    fn test_paren_expression() {
        assert_eq!(
            read_paren_expression("((2 + 4 * 9) * (6 + 9 * 8 + 6) + 6) + 2 + 4 * 2"),
            ("(2 + 4 * 9) * (6 + 9 * 8 + 6) + 6", " + 2 + 4 * 2")
        );
    }

    #[test]
    fn test_number() {
        assert_eq!(read_number("243 + 3"), ("243", " + 3"),);
    }

    #[test]
    fn test_parse() {
        assert_eq!(parse_expression("23"), vec![Expression::Num(23)]);
        assert_eq!(
            parse_expression("2 + 3 * 4"),
            vec![
                Expression::Num(2),
                Expression::OpAdd,
                Expression::Num(3),
                Expression::OpMul,
                Expression::Num(4)
            ]
        );
        assert_eq!(
            parse_expression("2 + (3 * 4)"),
            vec![
                Expression::Num(2),
                Expression::OpAdd,
                Expression::SubExpr(vec![
                    Expression::Num(3),
                    Expression::OpMul,
                    Expression::Num(4)
                ])
            ]
        );
    }

    #[test]
    fn test_basics1() {
        assert_eq!(eval_part1(&parse_expression("23")), 23);
        assert_eq!(eval_part1(&parse_expression("2 + 3 * 4")), 20);
        assert_eq!(eval_part1(&parse_expression("2 + (3 * 4)")), 14);
    }

    fn test_basics2() {
        assert_eq!(eval_part2(&parse_expression("23")), 23);
        assert_eq!(eval_part2(&parse_expression("2 + 3 * 2 * 3 + 2")), 50);
        assert_eq!(eval_part2(&parse_expression("2 + (3 * 4)")), 14);
    }

    #[test]
    fn test_problems1() {
        assert_eq!(eval_part1(&parse_expression("2 * 3 + (4 * 5)")), 26);
        assert_eq!(
            eval_part1(&parse_expression("5 + (8 * 3 + 9 + 3 * 4 * 3)")),
            437
        );
        assert_eq!(
            eval_part1(&parse_expression(
                "5 * 9 * (7 * 3 * 3 + 9 * 3 + (8 + 6 * 4))"
            )),
            12240
        );
        assert_eq!(
            eval_part1(&parse_expression(
                "((2 + 4 * 9) * (6 + 9 * 8 + 6) + 6) + 2 + 4 * 2"
            )),
            13632
        );
    }

    #[test]
    fn test_problems2() {
        assert_eq!(eval_part2(&parse_expression("2 * 3 + (4 * 5)")), 46);
        assert_eq!(
            eval_part2(&parse_expression("5 + (8 * 3 + 9 + 3 * 4 * 3)")),
            1445
        );
        assert_eq!(
            eval_part2(&parse_expression(
                "5 * 9 * (7 * 3 * 3 + 9 * 3 + (8 + 6 * 4))"
            )),
            669060
        );
        assert_eq!(
            eval_part2(&parse_expression(
                "((2 + 4 * 9) * (6 + 9 * 8 + 6) + 6) + 2 + 4 * 2"
            )),
            23340
        );
    }
}