use bevy::prelude::*;
use std::collections::{BTreeMap, BTreeSet};
use std::fs::File;
use std::io::{BufRead, BufReader};

fn main() {
    App::build()
        .add_resource(WindowDescriptor {
            title: "Advent of Code".to_string(),
            width: 1920,
            height: 1080,
            ..Default::default()
        })
        .add_resource(ClearColor(Color::rgb(0., 0., 0.)))
        .add_startup_system(setup_camera.system())
        .add_startup_system(read_input_file.system())
        .add_stage("calculate")
        .add_system_to_stage("calculate", calculate_anyone_groups.system())
        .add_system_to_stage("calculate", calculate_everyone_groups.system())
        .add_stage("report")
        .add_system_to_stage("report", count_anyone_group_sizes.system())
        .add_system_to_stage("report", count_everyone_group_sizes.system())
        //.add_plugins(DefaultPlugins)
        .run();
}

fn setup_camera(mut commands: Commands) {
    commands.spawn(Camera2dComponents::default());
}

struct Individual;
struct Anyone;
struct Everyone;

struct Group(usize);
struct Answers(BTreeSet<char>);

fn read_input_file(mut commands: Commands) {
    let f = File::open("./inputs/day_6.txt").unwrap();

    let mut current_group = 0;
    for line in BufReader::new(f).lines() {
        let line = line.unwrap();
        let line = line.trim();
        if line.is_empty() {
            current_group += 1;
        } else {
            let answers = line.chars().collect();
            commands.spawn((Individual, Group(current_group), Answers(answers)));
        }
    }
}

fn calculate_anyone_groups(
    mut commands: Commands,
    individuals: Query<With<Individual, (&Group, &Answers)>>,
) {
    let mut anyone_groups = BTreeMap::new();
    for (group_id, answer) in individuals.iter() {
        let updated = anyone_groups
            .entry(group_id.0)
            .or_insert(BTreeSet::new())
            .union(&answer.0)
            .cloned()
            .collect();
        anyone_groups.insert(group_id.0, updated);
    }

    for (group_id, anyone) in anyone_groups.into_iter() {
        commands.spawn((Anyone, Group(group_id), Answers(anyone)));
    }
}

fn calculate_everyone_groups(
    mut commands: Commands,
    individuals: Query<With<Individual, (&Group, &Answers)>>,
) {
    let mut anyone_groups = BTreeMap::new();
    for (group_id, answer) in individuals.iter() {
        let updated = anyone_groups
            .entry(group_id.0)
            .or_insert(answer.0.clone())
            .intersection(&answer.0)
            .cloned()
            .collect();
        anyone_groups.insert(group_id.0, updated);
    }

    for (group_id, anyone) in anyone_groups.into_iter() {
        commands.spawn((Everyone, Group(group_id), Answers(anyone)));
    }
}

fn count_anyone_group_sizes(groups: Query<With<Anyone, &Answers>>) {
    let sum: usize = groups.iter().map(|g| g.0.len()).sum();
    println!("Anyone: {}", sum);
}

fn count_everyone_group_sizes(groups: Query<With<Everyone, &Answers>>) {
    let sum: usize = groups.iter().map(|g| g.0.len()).sum();
    println!("Everyone: {}", sum);
}