use {
    anyhow::{
        anyhow,
        Context,
        Result,
    },
    std::{
        collections::BTreeMap,
        env::args,
        fs::File,
        io::{
            BufRead,
            BufReader,
        },
        path::Path,
    },
};

fn read_input<T: AsRef<Path>>(filename: T) -> Result<Vec<usize>> {
    let reader = BufReader::new(File::open(filename)?);

    reader.lines().next().ok_or(anyhow!(""))??.split(',').map(
        |s| {
            let n = s.parse()?;
            Ok(n)
        }
    ).collect()
}

fn part1<'a, I: Copy + IntoIterator<Item = &'a usize>>(input: I) -> Result<usize> {
    let mut subs: BTreeMap<usize, usize> = BTreeMap::new();
    let mut fuel = 0;
    for position in input {
        *subs.entry(*position).or_insert(0) += 1;
    }
    
    while subs.len() > 1 {
        let first_position = *subs.keys().next().ok_or(anyhow!("Couldn't get first position."))?;
        let last_position = *subs.keys().rev().next()
            .ok_or(anyhow!("Couldn't get last position."))?;
        let first_value = *subs.get(&first_position).ok_or(anyhow!("Couldn't get first value."))?;
        let last_value = *subs.get(&last_position).ok_or(anyhow!("Couldn't get last value."))?;
        if first_value <= last_value {
            let value = subs.remove(&first_position)
                .ok_or(anyhow!("Couldn't remove first value"))?;
            let new_position = first_position + 1;
            *subs.entry(new_position).or_insert(0) += value;
            fuel += value;
        }
        if first_value >= last_value {
            let value = subs.remove(&last_position).ok_or(anyhow!("Couldn't remove last value"))?;
            let new_position = last_position - 1;
            *subs.entry(new_position).or_insert(0) += value;
            fuel += value;
        }
    }
    Ok(fuel)
}

fn part2<'a, I: Copy + IntoIterator<Item = &'a usize>>(input: I) -> Result<usize> {
    #[derive(Debug)]
    struct Occupancy {
        cost: usize,
        vessel_count: usize,
    }

    impl Occupancy {
        fn new() -> Self {
            let cost = 0;
            let vessel_count = 0;
            Self {
                cost,
                vessel_count,
            }
        }

        fn add_vessel(&mut self) {
            self.cost += 1;
            self.vessel_count += 1;
        }

        fn merge(&mut self, other: Occupancy) -> Result<()> {
            self.cost += other.cost + other.vessel_count;
            self.vessel_count += other.vessel_count;
            Ok(())
        }
    }

    let mut subs: BTreeMap<usize, Occupancy> = BTreeMap::new();
    let mut fuel = 0;
    for position in input {
        subs.entry(*position).or_insert_with(Occupancy::new).add_vessel();
    }
    
    while subs.len() > 1 {
        let first_position = *subs.keys().next().ok_or(anyhow!("Couldn't get first position."))?;
        let last_position = *subs.keys().rev().next()
            .ok_or(anyhow!("Couldn't get last position."))?;
        let first_costs = (*subs.get(&first_position).ok_or(anyhow!("Couldn't get first value."))?)
            .cost;
        let first_fuel = first_costs;
        let last_costs = (*subs.get(&last_position).ok_or(anyhow!("Couldn't get last value."))?)
            .cost;
        let last_fuel = last_costs;
        if first_fuel <= last_fuel {
            let value = subs.remove(&first_position)
                .ok_or(anyhow!("Couldn't remove first value"))?;
            let new_position = first_position + 1;
            let costs = subs.entry(new_position).or_insert_with(Occupancy::new);
            costs.merge(value)?;
            fuel += first_fuel;
        }
        if first_fuel >= last_fuel {
            let value = subs.remove(&last_position).ok_or(anyhow!("Couldn't remove last value"))?;
            let new_position = last_position - 1;
            let costs = subs.entry(new_position).or_insert_with(Occupancy::new);
            costs.merge(value)?;
            fuel += last_fuel;
        }
    }
    Ok(fuel)
}

fn main() -> Result<()> {
    let ( do_part_1, do_part_2 ) = aoc::do_parts();

    let filename = args().nth(1).ok_or(anyhow!("Missing input filename"))?;
    let input = read_input(filename).context("Could not read input")?;
    if do_part_1 {
        let solution = part1(&input).context("No solution for part 1")?;
        println!("Part1, solution found to be: {}", solution);
    }
    if do_part_2 {
        let solution = part2(&input).context("No solution for part 2")?;
        println!("Part2, solution found to be: {}", solution);
    }
    Ok(())
}