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

#[derive(Clone,Copy,Debug)]
struct LanternFish {
    days: u8,
}

impl LanternFish {
    pub fn new(days: u8) -> Self {
        Self {
            days,
        }
    }

    fn step(&mut self) -> bool {
        if self.days == 0 {
            self.days = 6;
            true
        } else {
            self.days -= 1;
            false
        }
    }
}

impl Default for LanternFish {
    fn default() -> Self {
        Self::new(8)
    }
}

struct Population {
    count: Vec<usize>,
}

impl Population {
    fn new<'a, I: IntoIterator<Item = &'a LanternFish>>(input: I) -> Self {
        let mut count = vec![0; 9];

        for fish in input {
            count[fish.days as usize] += 1;
        }

        Self {
            count,
        }
    }

    fn step(&mut self) {
        let iter = &mut self.count.iter_mut();
        let spawn = iter.next().unwrap();

        let mut new_count: Vec<_> = iter.map(|v| *v).collect();
        new_count.push(0);
        new_count[6] += *spawn;
        new_count[8] = *spawn;
        self.count = new_count;
    }

    fn sum(&self) -> usize {
        self.count.iter().sum()
    }
}

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

    let line = reader.lines().next().ok_or(anyhow!("Could not read input"))??;
    line.split(',').map(|v| {
        let days = v.parse().map_err(|err| anyhow!("{}", err))?;
        Ok(LanternFish::new(days))
    }).collect()
}

fn part1<'a, I: IntoIterator<Item = &'a LanternFish>>(input: I) -> Result<usize> {
    let mut population: Vec<_> = input.into_iter().copied().collect();
    for _ in 0..80 {
        let mut children = vec![];
        for fish in &mut population {
            if fish.step() {
                children.push(LanternFish::default())
            }
        }
        population.append(&mut children);
    }
    Ok(population.len())
}

fn part2<'a, I: IntoIterator<Item = &'a LanternFish>>(input: I) -> Result<usize> {
    let mut population = Population::new(input);
    for _ in 0..256 {
        population.step();
    }

    Ok(population.sum())
}

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(())
}