use std::collections::HashSet;
use std::thread;

#[path = "../lib.rs"] mod lib;

fn get_blueprints() -> Vec<[[usize; 4]; 4]> {
    let mut blueprints = vec![];
    for line in lib::iter_input() {
        let parts = line.split(" ").collect::<Vec<&str>>();
        let ore_robot_ore_cost = parts[6].parse().unwrap();
        let clay_robot_ore_cost = parts[12].parse().unwrap();
        let obsidian_robot_ore_cost = parts[18].parse().unwrap();
        let obsidian_robot_clay_cost = parts[21].parse().unwrap();
        let geode_robot_ore_cost = parts[27].parse().unwrap();
        let geode_robot_obsidian_cost = parts[30].parse().unwrap();
        blueprints.push([
            [ore_robot_ore_cost, 0, 0, 0],
            [clay_robot_ore_cost, 0, 0, 0],
            [obsidian_robot_ore_cost, obsidian_robot_clay_cost, 0, 0],
            [geode_robot_ore_cost, 0, geode_robot_obsidian_cost, 0],
        ]);
    }
    return blueprints;
}

fn get_baseline(blueprint: &[[usize; 4]; 4], turns: usize) -> usize {
    // thought up strategy that should yield good results.
    // I am not sure if this is ideal though, so we will do a thorough search afterwards.

    let mut n = turns;
    let mut resources = [0, 0, 0, 0];
    let mut robots = [1, 0, 0, 0];

    // build enough ore robots to sustain geode robot production
    while n > 0 && robots[0] < blueprint[3][0] {
        let build = (0..4).all(|j| resources[j] >= blueprint[0][j]);
        n -= 1;
        for j in 0..4 {
            resources[j] += robots[j];
        }
        if build {
            for j in 0..4 {
                resources[j] -= blueprint[0][j];
            }
            robots[0] += 1;
        }
    }

    // build enough obsidian robots to sustain geode robot production
    // and then only build geode robots
    while n > 0 {
        let can_clay = (0..4).all(|j| resources[j] >= blueprint[1][j]);
        let can_obsidian = (0..4).all(|j| resources[j] >= blueprint[2][j]);
        let can_geode = (0..4).all(|j| resources[j] >= blueprint[3][j]);

        let need_obsidian = robots[2] < blueprint[3][2];
        let need_clay = need_obsidian && robots[1] < blueprint[2][1] && resources[1] < blueprint[2][1];

        n -= 1;
        for j in 0..4 {
            resources[j] += robots[j];
        }

        if can_geode {
            for j in 0..4 {
                resources[j] -= blueprint[3][j];
            }
            robots[3] += 1;
        } else if need_obsidian && can_obsidian {
            for j in 0..4 {
                resources[j] -= blueprint[2][j];
            }
            robots[2] += 1;
        } else if need_clay && can_clay {
            for j in 0..4 {
                resources[j] -= blueprint[1][j];
            }
            robots[1] += 1;
        }
    }

    return resources[3];
}

fn get_max_geodes(blueprint: &[[usize; 4]; 4], turns: usize) -> usize {
    let mut best_geodes = get_baseline(blueprint, turns);
    let mut queue = vec![
        ([0, 0, 0, 0], [1, 0, 0, 0], turns),
    ];
    let mut used = HashSet::new();

    while let Some((resources, robots, n)) = queue.pop() {
        let min_geodes = resources[3] + robots[3] * n;
        let potential = n * (n - 1) / 2;
        if min_geodes + potential <= best_geodes {
            continue;
        }

        if used.contains(&(resources, robots, n)) {
            continue;
        } else {
            used.insert((resources, robots, n));
        }

        let mut build = vec![];

        // if it is possible to build a geode, always do that
        if (0..4).all(|j| resources[j] >= blueprint[3][j]) {
            build.push(3);
        } else {
            let need_obsidian = robots[2] < blueprint[3][2];
            let need_clay = need_obsidian && robots[1] < blueprint[2][1];
            let need_ore = robots[0] < blueprint[3][0] || (
                need_obsidian && robots[0] < blueprint[2][0]
            ) || (
                need_clay && robots[0] < blueprint[1][0]
            );

            build.push(4);
            if need_ore && (0..4).all(|j| resources[j] >= blueprint[0][j]) {
                build.push(0);
            }
            if need_clay && (0..4).all(|j| resources[j] >= blueprint[1][j]) {
                build.push(1);
            }
            if need_obsidian && (0..4).all(|j| resources[j] >= blueprint[2][j]) {
                build.push(2);
            }
        }

        for i in build {
            let mut next_resources = [
                resources[0] + robots[0],
                resources[1] + robots[1],
                resources[2] + robots[2],
                resources[3] + robots[3],
            ];
            let mut next_robots = robots.clone();
            if i < 4 {
                for j in 0..4 {
                    next_resources[j] -= blueprint[i][j];
                }
                next_robots[i] += 1;
            }

            let min_geodes = next_resources[3] + next_robots[3] * (n - 1);
            if min_geodes > best_geodes {
                best_geodes = min_geodes;
            }

            queue.push((next_resources, next_robots, n - 1));
        }
    }

    return best_geodes;
}

fn part1(blueprints: &Vec<[[usize; 4]; 4]>) -> usize {
    let mut handles = vec![];
    for (i, blueprint) in blueprints.iter().enumerate() {
        let b = blueprint.clone();
        handles.push(thread::spawn(move || {
            return get_max_geodes(&b, 24) * (i + 1);
        }));
    }
    return handles.into_iter()
        .map(|handle| handle.join().unwrap())
        .sum();
}

fn part2(blueprints: &Vec<[[usize; 4]; 4]>) -> usize {
    let mut handles = vec![];
    for blueprint in blueprints.iter().take(3) {
        let b = blueprint.clone();
        handles.push(thread::spawn(move || {
            return get_max_geodes(&b, 32);
        }));
    }
    return handles.into_iter()
        .map(|handle| handle.join().unwrap())
        .product();
}

fn main() {
    let blueprints = get_blueprints();
    println!("part1: {}", part1(&blueprints));
    println!("part2: {}", part2(&blueprints));
}