use derive_more;
use std::io;
use std::io::prelude::*;
use std::iter;
use std::process;

use structopt::StructOpt;

#[derive(Debug, StructOpt)]
#[structopt(name = "Day 1: The Tyranny of the Rocket Equation")]
/// Calculates the fuel needed for your rocket to save Santa.
///
/// The weight of each module is read from stdin, one module weight
/// per line. See https://adventofcode.com/2019/day/1 for details.
struct Opt {
    /// Includes the weight of fuel
    #[structopt(short = "i", long = "include-fuel-weight")]
    include_fuel_weight: bool,
}

fn main() {
    let stdin = io::stdin();
    let opt = Opt::from_args();

    let input = stdin.lock().lines().map(|l| match l {
        Ok(s) => match s.parse::<Module>() {
            Ok(module) => module,
            Err(e) => {
                eprintln!("Invalid input \"{}\": {}", s, e);
                process::exit(1);
            }
        },
        Err(e) => {
            eprintln!("Error reading input: {}", e);
            process::exit(1);
        }
    });

    println!("{}", fuel_required(input, opt.include_fuel_weight))
}

fn fuel_required(it: impl Iterator<Item = Module>, include_fuel_weight: bool) -> Fuel {
    it.map(if include_fuel_weight {
        Module::fuel_including_fuel_weight
    } else {
        Module::fuel_excluding_fuel_weight
    })
    .sum()
}

#[derive(Debug, derive_more::FromStr, Clone, Copy)]
struct Module {
    weight: Weight,
}

impl Module {
    fn fuel_excluding_fuel_weight(self) -> Fuel {
        self.weight.required_fuel()
    }

    fn fuel_including_fuel_weight(self) -> Fuel {
        iter::successors(Some(self.weight.required_fuel()), |fuel| {
            if fuel.is_zero() {
                None
            } else {
                Some(fuel.weight().required_fuel())
            }
        })
        .sum()
    }
}

#[derive(Debug, derive_more::FromStr, Clone, Copy)]
struct Weight(u32);

impl Weight {
    fn required_fuel(self) -> Fuel {
        Fuel((self.0 / 3).saturating_sub(2))
    }
}

#[derive(Debug, derive_more::Add, derive_more::Sum, Clone, Copy, derive_more::Display)]
struct Fuel(u32);

impl Fuel {
    fn weight(self) -> Weight {
        Weight(self.0)
    }

    fn is_zero(self) -> bool {
        self.0 == 0
    }
}