use nom::{
    character::complete::{digit1, line_ending},
    combinator::{map, map_res},
    multi::separated_list1,
    IResult,
};
use std::fs;

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

    dbg!(diagnostics.gamma());
    dbg!(diagnostics.epsilon());
    dbg!(diagnostics.oxygen());
    dbg!(diagnostics.co2());
    dbg!(diagnostics.gamma() * diagnostics.epsilon());
    dbg!(diagnostics.oxygen() * diagnostics.co2());

    Ok(())
}

#[derive(Debug, Clone)]
struct Diagnostics {
    bitsets: Vec<Bitset>,
    bitset_len: usize,
}

impl Diagnostics {
    fn new(bitsets: Vec<Bitset>) -> Diagnostics {
        Diagnostics {
            bitset_len: bitsets.iter().map(|b| b.len).max().unwrap_or(0),
            bitsets,
        }
    }

    fn gamma(&self) -> Bitset {
        let mut gamma = Bitset {
            bits: 0,
            len: self.bitset_len,
        };
        for bit in 0..gamma.len {
            let ones_count = self
                .bitsets
                .iter()
                .filter(|bitset| bitset.check_bit(bit))
                .count();
            if ones_count >= self.bitsets.len() / 2 {
                gamma.set_bit(bit);
            }
        }
        gamma
    }

    fn epsilon(&self) -> Bitset {
        !self.gamma()
    }

    fn oxygen(&self) -> Bitset {
        self.bit_criteria_match(|d| d.gamma())
    }

    fn co2(&self) -> Bitset {
        self.bit_criteria_match(|d| d.epsilon())
    }

    fn bit_criteria_match(&self, criteria: impl Fn(&Diagnostics) -> Bitset) -> Bitset {
        let mut candidates = self.clone();
        for bit in (0..self.bitset_len).rev() {
            let bit_criteria = criteria(&candidates).check_bit(bit);
            candidates
                .bitsets
                .retain(|candidate| candidate.check_bit(bit) == bit_criteria);
            if candidates.bitsets.len() == 1 {
                return candidates.bitsets[0];
            }
        }
        Bitset::default()
    }
}

#[derive(Default, Debug, Clone, Copy)]
struct Bitset {
    bits: u32,
    len: usize,
}

impl Bitset {
    fn check_bit(&self, bit: usize) -> bool {
        self.bits & (1 << bit) != 0
    }

    fn set_bit(&mut self, bit: usize) {
        self.bits |= 1 << bit;
    }
}

impl std::ops::Mul for Bitset {
    type Output = Self;
    fn mul(self, rhs: Bitset) -> Self::Output {
        Bitset {
            bits: self.bits * rhs.bits,
            len: self.len.max(rhs.len), // dodgy. This might need to grow.
        }
    }
}
impl std::ops::Not for Bitset {
    type Output = Bitset;
    fn not(self) -> <Self as std::ops::Not>::Output {
        Bitset {
            bits: !self.bits & ((1 << self.len) - 1),
            len: self.len,
        }
    }
}

fn parse_diagnostics(input: &str) -> IResult<&str, Diagnostics> {
    map(separated_list1(line_ending, parse_bitset), Diagnostics::new)(input)
}

fn parse_bitset(input: &str) -> IResult<&str, Bitset> {
    map_res(digit1, |num| {
        u32::from_str_radix(num, 2).map(|bits| Bitset {
            bits,
            len: num.len(),
        })
    })(input)
}