use aoc_runner_derive::{aoc, aoc_generator};
use std::collections::VecDeque;

#[derive(Debug, Eq, PartialEq, Copy, Clone)]
pub enum Token {
    Num(i64),
    Add,
    Mul,
    LParen,
    RParen,
}

impl Token {
    fn precedence(&self, part_2: bool) -> i64 {
        match self {
            Token::Add => {
                if part_2 {
                    1
                } else {
                    0
                }
            }
            Token::Mul => 0,
            _ => unimplemented!(),
        }
    }
}
#[derive(Debug, Eq, PartialEq)]
struct RPN(VecDeque<Token>);

impl RPN {
    fn from_tokens(tokens: &[Token], part_2: bool) -> Self {
        let mut operators: Vec<Token> = Vec::new();
        let mut output: VecDeque<Token> = VecDeque::new();
        for token in tokens.iter().cloned() {
            match &token {
                Token::Num(_) => output.push_back(token),
                Token::Add | Token::Mul => {
                    if let Some(top) = operators.last() {
                        if (*top != Token::LParen)
                            && top.precedence(part_2) >= token.precedence(part_2)
                        {
                            output.push_back(operators.pop().unwrap())
                        }
                    }
                    operators.push(token);
                }
                Token::LParen => operators.push(token),
                Token::RParen => {
                    while Some(&Token::LParen) != operators.last() {
                        output.push_back(operators.pop().expect("Unmatched parenthesis!"));
                    }
                    if let Some(Token::LParen) = operators.last() {
                        operators.pop();
                    }
                }
            }
        }
        while !operators.is_empty() {
            output.push_back(operators.pop().unwrap());
        }
        RPN(output)
    }

    fn eval(&self) -> i64 {
        let mut stack = vec![];
        for token in &self.0 {
            match token {
                Token::Num(n) => {
                    stack.push(*n);
                }
                Token::Add => {
                    let a = stack.pop().unwrap();
                    let b = stack.pop().unwrap();
                    stack.push(a + b);
                }
                Token::Mul => {
                    let a = stack.pop().unwrap();
                    let b = stack.pop().unwrap();
                    stack.push(a * b);
                }
                _ => (),
            }
        }
        if stack.len() != 1 {
            panic!("Error evaluating, leftover stack: {:?}", stack);
        }
        return stack.pop().unwrap();
    }
}

impl Token {
    fn lex(s: &str) -> Token {
        match s {
            "+" => Token::Add,
            "*" => Token::Mul,
            "(" => Token::LParen,
            ")" => Token::RParen,
            n if n.chars().all(|c| c.is_numeric()) => Token::Num(n.parse().unwrap()),
            other => panic!("Failed to parse: {:?}", other),
        }
    }
}

#[aoc_generator(day18)]
pub fn input_generator(input: &str) -> Vec<Vec<Token>> {
    let exprs: Vec<Vec<Token>> = input
        .lines()
        .map(|l| {
            l.replace("(", "( ")
                .replace(")", " )")
                .split_ascii_whitespace()
                .map(|s| Token::lex(s))
                .collect()
        })
        .collect();
    exprs
}

#[aoc(day18, part1)]
pub fn solve_part1(input: &[Vec<Token>]) -> usize {
    input
        .iter()
        .map(|e| RPN::from_tokens(e, false).eval())
        .sum::<i64>() as usize
}

#[aoc(day18, part2)]
pub fn solve_part2(input: &[Vec<Token>]) -> usize {
    input
        .iter()
        .map(|e| RPN::from_tokens(e, true).eval())
        .sum::<i64>() as usize
}

#[cfg(test)]
mod test {

    #[test]
    fn part1() {
        use super::{input_generator, solve_part1};
        for (input, res) in &[
            ("2 * 3 + (4 * 5)", 26),
            ("1 + (2 * 3) + (4 * (5 + 6))", 51),
            ("5 + (8 * 3 + 9 + 3 * 4 * 3)", 437),
            ("5 * 9 * (7 * 3 * 3 + 9 * 3 + (8 + 6 * 4))", 12240),
            ("((2 + 4 * 9) * (6 + 9 * 8 + 6) + 6) + 2 + 4 * 2", 13632),
        ] {
            assert_eq!(solve_part1(&input_generator(input)), *res, "{}", input);
        }
    }

    #[test]
    fn part2() {
        use super::{input_generator, solve_part2};
        for (input, res) in &[
            ("2 * 3 + (4 * 5)", 46),
            ("1 + (2 * 3) + (4 * (5 + 6))", 51),
            ("5 + (8 * 3 + 9 + 3 * 4 * 3)", 1445),
            ("5 * 9 * (7 * 3 * 3 + 9 * 3 + (8 + 6 * 4))", 669060),
            ("((2 + 4 * 9) * (6 + 9 * 8 + 6) + 6) + 2 + 4 * 2", 23340),
        ] {
            assert_eq!(solve_part2(&input_generator(input)), *res, "{}", input);
        }
    }
}