use nom::{
    branch::alt,
    bytes::complete::tag,
    character::complete::{line_ending, u32},
    combinator::map,
    multi::separated_list1,
    sequence::tuple,
    IResult,
};
use std::{cell::RefCell, collections::BTreeMap, fs};

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let input = fs::read_to_string("inputs/day_19.txt")?;
    let blueprints = Blueprints::parser(&input).unwrap().1;

    dbg!(blueprints.quality_sum());

    let part_2_state = State {
        time_remaining: 32,
        resources: Resources::default(),
        robots: Resources {
            ore: 1,
            ..Resources::default()
        },
    };

    dbg!(blueprints
        .0
        .iter()
        .take(3)
        .map(|b| {
            let result = b.max_geodes(part_2_state.clone());
            b.clear_cache();
            result
        })
        .product::<u32>());

    Ok(())
}

#[derive(Debug, Clone)]
struct Blueprints(Vec<Blueprint>);

#[derive(Debug, Clone)]
struct Blueprint {
    id: u32,
    ore_robot: Resources,
    clay_robot: Resources,
    obsidian_robot: Resources,
    geode_robot: Resources,
    max_geode_cache: RefCell<BTreeMap<State, u32>>,
}

#[derive(Debug, Default, Clone, PartialEq, Eq, PartialOrd, Ord)]
struct Resources {
    ore: u32,
    clay: u32,
    obsidian: u32,
    geodes: u32,
}

#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
struct State {
    time_remaining: u32,
    resources: Resources,
    robots: Resources,
}

impl Blueprints {
    fn parser(input: &str) -> IResult<&str, Self> {
        map(separated_list1(line_ending, Blueprint::parser), Blueprints)(input)
    }

    fn quality_sum(&self) -> u32 {
        self.0.iter().map(|b| b.quality_level()).sum()
    }
}

impl Blueprint {
    fn parser(input: &str) -> IResult<&str, Self> {
        map(
            tuple((
                tag("Blueprint "),
                u32,
                tag(": Each ore robot costs "),
                Resources::parser,
                tag(". Each clay robot costs "),
                Resources::parser,
                tag(". Each obsidian robot costs "),
                Resources::parser,
                tag(". Each geode robot costs "),
                Resources::parser,
                tag("."),
            )),
            |(_, id, _, ore_robot, _, clay_robot, _, obsidian_robot, _, geode_robot, _)| {
                Blueprint {
                    id,
                    ore_robot,
                    clay_robot,
                    obsidian_robot,
                    geode_robot,
                    max_geode_cache: RefCell::new(BTreeMap::new()),
                }
            },
        )(input)
    }

    fn quality_level(&self) -> u32 {
        let result = self.id
            * self.max_geodes(State {
                time_remaining: 24,
                resources: Resources::default(),
                robots: Resources {
                    ore: 1,
                    ..Resources::default()
                },
            });
        self.clear_cache();
        result
    }

    fn max_geodes(&self, state: State) -> u32 {
        if state.time_remaining == 0 {
            state.resources.geodes
        } else {
            if let Some(cache) = self.max_geode_cache.borrow().get(&state) {
                return cache.clone();
            }

            let calculated = state
                .possible_next_states(&self)
                .into_iter()
                .map(|next_state| self.max_geodes(next_state))
                .max()
                .unwrap();
            self.max_geode_cache
                .borrow_mut()
                .insert(state.clone(), calculated);
            calculated
        }
    }

    fn clear_cache(&self) {
        self.max_geode_cache.borrow_mut().clear();
    }
}

impl Resources {
    fn parser(input: &str) -> IResult<&str, Self> {
        map(
            separated_list1(
                tag(" and "),
                tuple((
                    u32,
                    tag(" "),
                    alt((tag("ore"), tag("clay"), tag("obsidian"))),
                )),
            ),
            |resources| {
                let mut ore = 0;
                let mut clay = 0;
                let mut obsidian = 0;

                for (amount, _, resource) in resources {
                    match resource {
                        "ore" => ore = amount,
                        "clay" => clay = amount,
                        "obsidian" => obsidian = amount,
                        unknown => panic!("{} isn't one of the specified tags", unknown),
                    }
                }

                Resources {
                    ore,
                    clay,
                    obsidian,
                    geodes: 0,
                }
            },
        )(input)
    }
}

impl State {
    fn possible_next_states(&self, blueprint: &Blueprint) -> Vec<State> {
        let mut next_states = Vec::new();

        // build an ore robot
        if let Some(time) = self.time_to_afford(&blueprint.ore_robot) {
            let mut next_state = self.simulate_time_passing(time + 1);
            next_state.robots.ore += 1;
            next_state.resources -= &blueprint.ore_robot;
            next_states.push(next_state);
        }

        // build a clay robot
        if let Some(time) = self.time_to_afford(&blueprint.clay_robot) {
            let mut next_state = self.simulate_time_passing(time + 1);
            next_state.robots.clay += 1;
            next_state.resources -= &blueprint.clay_robot;
            next_states.push(next_state);
        }

        // build an obsidian robot
        if let Some(time) = self.time_to_afford(&blueprint.obsidian_robot) {
            let mut next_state = self.simulate_time_passing(time + 1);
            next_state.robots.obsidian += 1;
            next_state.resources -= &blueprint.obsidian_robot;
            next_states.push(next_state);
        }

        // build a geode robot
        if let Some(time) = self.time_to_afford(&blueprint.geode_robot) {
            let mut next_state = self.simulate_time_passing(time + 1);
            next_state.robots.geodes += 1;
            next_state.resources -= &blueprint.geode_robot;
            next_states.push(next_state);
        }

        // or just do nothing
        if next_states.len() == 0 {
            next_states.push(self.simulate_time_passing(self.time_remaining));
        }

        next_states
    }

    fn simulate_time_passing(&self, time: u32) -> State {
        State {
            time_remaining: self.time_remaining - time,
            resources: &self.resources + &(&self.robots * time),
            robots: self.robots.clone(),
        }
    }

    fn time_to_afford(&self, cost: &Resources) -> Option<u32> {
        let mut time_required = 0;

        if self.resources.ore < cost.ore {
            if self.robots.ore == 0 {
                return None;
            }
            time_required = time_required
                .max((cost.ore - self.resources.ore + self.robots.ore - 1) / self.robots.ore);
        }

        if self.resources.clay < cost.clay {
            if self.robots.clay == 0 {
                return None;
            }
            time_required = time_required
                .max((cost.clay - self.resources.clay + self.robots.clay - 1) / self.robots.clay);
        }

        if self.resources.obsidian < cost.obsidian {
            if self.robots.obsidian == 0 {
                return None;
            }
            time_required = time_required.max(
                (cost.obsidian - self.resources.obsidian + self.robots.obsidian - 1)
                    / self.robots.obsidian,
            );
        }

        if self.resources.geodes < cost.geodes {
            if self.robots.geodes == 0 {
                return None;
            }
            time_required = time_required.max(
                (cost.geodes - self.resources.geodes + self.robots.geodes - 1) / self.robots.geodes,
            );
        }

        if time_required + 1 >= self.time_remaining {
            None
        } else {
            Some(time_required)
        }
    }
}

impl ::core::ops::Add<&Resources> for &Resources {
    type Output = Resources;
    fn add(self, rhs: &Resources) -> Resources {
        Resources {
            ore: self.ore.add(rhs.ore),
            clay: self.clay.add(rhs.clay),
            obsidian: self.obsidian.add(rhs.obsidian),
            geodes: self.geodes.add(rhs.geodes),
        }
    }
}

impl ::core::ops::Sub<&Resources> for &Resources {
    type Output = Resources;
    fn sub(self, rhs: &Resources) -> Resources {
        Resources {
            ore: self.ore.sub(rhs.ore),
            clay: self.clay.sub(rhs.clay),
            obsidian: self.obsidian.sub(rhs.obsidian),
            geodes: self.geodes.sub(rhs.geodes),
        }
    }
}

impl ::core::ops::SubAssign<&Resources> for Resources {
    fn sub_assign(&mut self, rhs: &Resources) {
        self.ore -= rhs.ore;
        self.clay -= rhs.clay;
        self.obsidian -= rhs.obsidian;
        self.geodes -= rhs.geodes;
    }
}

impl ::core::ops::Mul<u32> for &Resources {
    type Output = Resources;
    fn mul(self, rhs: u32) -> Resources {
        Resources {
            ore: self.ore.mul(rhs),
            clay: self.clay.mul(rhs),
            obsidian: self.obsidian.mul(rhs),
            geodes: self.geodes.mul(rhs),
        }
    }
}

#[test]
fn time_to_afford() {
    let current_state = State {
        time_remaining: 17,
        resources: Resources {
            ore: 1,
            clay: 6,
            obsidian: 0,
            geodes: 0,
        },
        robots: Resources {
            ore: 1,
            clay: 3,
            obsidian: 0,
            geodes: 0,
        },
    };

    let time_to_afford = current_state.time_to_afford(&Resources {
        ore: 3,
        clay: 14,
        obsidian: 0,
        geodes: 0,
    });
    assert_eq!(time_to_afford, Some(3));
}