use std::collections::HashSet;

fn decode(inputs: &[HashSet<char>], element: &HashSet<char>) -> usize {
    let length = element.len();
    if length == inputs[0].len() {
        return 1;
    } else if length == inputs[1].len() {
        return 7;
    } else if length == inputs[2].len() {
        return 4;
    } else if length == inputs[9].len() {
        return 8;
    }

    if length == 5 {
        if inputs[0].difference(element).count() == 0 {
            return 3;
        } else if inputs[2].difference(element).count() == 1 {
            return 5;
        } else {
            return 2;
        }
    }
    if inputs[0].difference(element).count() > 0 {
        return 6;
    } else if inputs[2].difference(element).count() > 0 {
        return 0;
    } else {
        return 9;
    }
}

fn aoc_08_part02(input: &mut [(Vec<HashSet<char>>, Vec<HashSet<char>>)]) -> usize {
    let mut result = 0;

    for (first, last) in input.iter_mut() {
        // Sorting the input makes it WAY more easier...
        first.sort_unstable_by_key(|l| l.len());

        // Now we try to decode...
        result += last.iter().map(|element| decode(&first, element))
            .fold(0, |acc, val| acc * 10 + val);
    }

    result
}

fn aoc_08_part01(input: &[(Vec<HashSet<char>>, Vec<HashSet<char>>)]) -> usize {
    input.iter().map(|(_first, last)| {
        last.iter().filter_map(|element| {
            let l = element.len();
            if l == 2 || l == 3 || l == 4 || l == 7 {
                return Some(l);
            }
            None
        }).count()
    }).sum()
}

fn main() {
    let mut s = include_str!("../input.txt").lines().map(|l| {
        let before_after = l.split_once(" | ").unwrap_or(("", ""));
        (
            before_after.0.split_whitespace().map(|d| d.chars().collect::<HashSet<_>>()).collect(),
            before_after.1.split_whitespace().map(|d| d.chars().collect::<HashSet<_>>()).collect()
        )
    }).collect::<Vec<(Vec<HashSet<char>>, Vec<HashSet<char>>)>>();

    let result = aoc_08_part01(&s);
    println!("Result {}", result);

    let result = aoc_08_part02(&mut s);
    println!("Result {}", result);
}