use itertools::Itertools;
use std::collections::HashMap;

use id_arena::{Arena, Id};

type MonkeyNumber = i64;
type NodeId = Id<Node>;

#[derive(Debug)]
enum Node {
    Uninitialized,
    Unit(MonkeyNumber),
    Add(NodeId, NodeId),
    Sub(NodeId, NodeId),
    Mul(NodeId, NodeId),
    Div(NodeId, NodeId),
}

impl Node {
    /// Evaluates the expression tree with the given Node as root node and returns the calculated result.
    ///
    /// Uses recursion under-the-hood to implement the evaluation.
    /// The recursive strategy is essentially a depth-first-search algorithm.
    fn evaluate(&self, nodes: &Arena<Node>) -> MonkeyNumber {
        match self {
            Node::Unit(value) => *value,
            Node::Add(lhs, rhs) => nodes[*lhs].evaluate(nodes) + nodes[*rhs].evaluate(nodes),
            Node::Sub(lhs, rhs) => nodes[*lhs].evaluate(nodes) - nodes[*rhs].evaluate(nodes),
            Node::Mul(lhs, rhs) => nodes[*lhs].evaluate(nodes) * nodes[*rhs].evaluate(nodes),
            Node::Div(lhs, rhs) => nodes[*lhs].evaluate(nodes) / nodes[*rhs].evaluate(nodes),
            Node::Uninitialized => panic!("encountered uninitialized Node during evaluation"),
        }
    }
}

fn main() {
    // Use command line arguments to specify the input filename.
    let args: Vec<String> = std::env::args().collect();
    if args.len() < 2 {
        panic!("Usage: ./main <input-file>\nNo input file provided. Exiting.");
    }

    // Next, read the contents of the input file into a string for easier processing.
    let input = std::fs::read_to_string(&args[1]).expect("Error opening file");

    // --- TASK BEGIN ---

    // The complete collection of operation nodes.
    let mut nodes = Arena::<Node>::with_capacity(input.lines().count());
    // Dictionary mapping the 4-letter identifiers to their respective NodeId.
    let mut identifiers = HashMap::<[u8; 4], NodeId>::new();

    // First, create unitinitialized nodes and populate the identifier-dictionary.
    for line in input.lines() {
        let node = nodes.alloc(Node::Uninitialized);
        identifiers.insert(line[0..4].as_bytes().try_into().unwrap(), node);
    }

    // Then, actually populate the nodes.
    for line in input.lines() {
        // Split into node_id and expression.
        let (node_id, expr) = line.split(": ").collect_tuple().unwrap();

        // Convert node_id from `&str` to `&[u8; 4]`.
        let node_id: &[u8; 4] = node_id.as_bytes().try_into().unwrap();
        // Grab the actual NodeId with the identifier.
        let node = identifiers[node_id];

        // Try to see if we can parse
        // -> an operator character (+, -, *, /)
        // -> a left- and right-hand-side of the operation
        // -> an integer for the whole expression
        let op_char = expr.chars().nth(5);
        let lhs: Option<&[u8; 4]> = expr
            .as_bytes()
            .get(0..4)
            .map(|e| e.try_into().ok())
            .flatten();
        let rhs: Option<&[u8; 4]> = expr
            .as_bytes()
            .get(7..11)
            .map(|e| e.try_into().ok())
            .flatten();
        let monkey_number = expr.parse::<MonkeyNumber>().ok();

        // Now actually assign the node, depending on what we could and could not match.
        nodes[node] = match (op_char, lhs, rhs, monkey_number) {
            (Some('+'), Some(lhs), Some(rhs), None) => {
                Node::Add(identifiers[lhs], identifiers[rhs])
            }
            (Some('-'), Some(lhs), Some(rhs), None) => {
                Node::Sub(identifiers[lhs], identifiers[rhs])
            }
            (Some('*'), Some(lhs), Some(rhs), None) => {
                Node::Mul(identifiers[lhs], identifiers[rhs])
            }
            (Some('/'), Some(lhs), Some(rhs), None) => {
                Node::Div(identifiers[lhs], identifiers[rhs])
            }
            (_, _, _, Some(monkey_number)) => Node::Unit(monkey_number),
            _ => panic!("could not parse line in input"),
        };
    }

    // Evaluate the expression tree and print the result.
    println!("Result: {}", nodes[identifiers[b"root"]].evaluate(&nodes));
}