use crate::utils;

pub fn both() {
    let input: Vec<usize> = utils::read_file("input/day08.txt")
        .split(' ')
        .map(|x| x.parse().unwrap())
        .collect();

    let (tree, _) = Node::extract_from(&input);
    println!("Metadata sum: {}", tree.metadata_sum());
    println!("Task2 sum: {}", tree.task2());
}

struct Node {
    metadata: Vec<usize>,
    children: Vec<Node>,
}

impl Node {
    fn extract_from(input: &[usize]) -> (Self, &[usize]) {
        let mut rest_input = &input[2..];
        let mut num_children: usize = input[0];
        let num_meta: usize = input[1];
        let mut children = Vec::new();
        while num_children > 0 {
            let (node, rest) = Node::extract_from(rest_input);
            children.push(node);
            rest_input = rest;
            num_children -= 1;
        }

        let metadata = rest_input[..num_meta].to_vec();
        (Node { metadata, children }, &rest_input[num_meta..])
    }

    fn metadata_sum(&self) -> usize {
        self.metadata.iter().sum::<usize>()
            + self
                .children
                .iter()
                .map(|child| {
                    let s: usize = child.metadata_sum();
                    s
                })
                .sum::<usize>()
    }

    fn task2(&self) -> usize {
        if !self.children.is_empty() {
            self.metadata
                .iter()
                .map(|meta| match self.children.get(*meta - 1) {
                    Some(child) => child.task2(),
                    None => 0,
                })
                .sum()
        } else {
            self.metadata_sum()
        }
    }
}