use aoc_runner_derive::{aoc, aoc_generator};
use std::collections::HashMap;

#[derive(Debug, Clone)]
pub enum Command {
    Mask { and: u64, or: u64, raw: Vec<char> },
    Write { val: u64, addr: u64 },
}

#[inline(always)]
fn set_bit(v: u64, bit: usize, val: bool) -> u64 {
    let mut v = v;
    let bv = 1 << bit;
    if val {
        v |= bv;
    } else {
        v &= !bv;
    }
    return v;
}

impl Command {
    fn apply_mask_v1(&self, val: u64) -> u64 {
        if let Command::Mask { and, or, .. } = self {
            return (val & and) | or;
        };
        panic!("Can't apply non-mask command");
    }

    fn apply_mask_v2_inner(&self, addr: &mut u64, s: &[char], ret: &mut Vec<u64>) {
        let mut s_c = s.to_vec();
        for (idx, c) in s.iter().enumerate() {
            match *c {
                '0' => (),
                '1' => {
                    *addr = set_bit(*addr, idx, true);
                }
                'X' => {
                    s_c[idx] = '0';
                    *addr = set_bit(*addr, idx, true);
                    self.apply_mask_v2_inner(addr, &s_c, ret);
                    *addr = set_bit(*addr, idx, false);
                    self.apply_mask_v2_inner(addr, &s_c, ret);
                }
                _ => unreachable!(),
            }
        }
        ret.push(*addr);
    }

    fn apply_mask_v2(&self, addr: u64) -> Vec<u64> {
        let mut addr = addr;
        if let Command::Mask { raw, .. } = self {
            let mut raw = raw.clone();
            raw.reverse();
            let mut ret = vec![];
            self.apply_mask_v2_inner(&mut addr, &raw, &mut ret);
            ret.sort_unstable();
            ret.dedup();
            return ret;
        };
        panic!("Can't apply non-mask command");
    }
}

type Data = Vec<Command>;

#[aoc_generator(day14)]
pub fn input_generator(input: &str) -> Data {
    let mut ret = vec![];
    for line in input.lines() {
        if line.starts_with("mask") {
            let mut mask_and = 0;
            let mut mask_or = 0;
            let mut char_buf = vec![];
            for (_, c) in line.split('=').nth(1).unwrap().trim().chars().enumerate() {
                char_buf.push(c);
                mask_and <<= 1;
                mask_or <<= 1;
                match c {
                    'X' => {
                        mask_and |= 1;
                        mask_or |= 0;
                    }
                    '1' => {
                        mask_and |= 1;
                        mask_or |= 1;
                    }
                    '0' => {
                        mask_and |= 0;
                        mask_or |= 0;
                    }
                    c => {
                        panic!("Invalid char in mask: {}", c);
                    }
                }
            }
            ret.push(Command::Mask {
                and: mask_and,
                or: mask_or,
                raw: char_buf,
            });
            continue;
        }
        if line.starts_with("mem") {
            let addr: u64 = line
                .chars()
                .skip_while(|&c| c != '[')
                .skip(1)
                .take_while(|&c| c != ']')
                .collect::<String>()
                .parse()
                .unwrap();
            let val: u64 = line
                .chars()
                .skip_while(|&c| c != '=')
                .skip(1)
                .skip_while(|&c| c.is_whitespace())
                .collect::<String>()
                .parse()
                .unwrap();
            ret.push(Command::Write { addr, val });
        }
    }
    ret
}

#[aoc(day14, part1)]
pub fn solve_part1(input: &[Command]) -> u64 {
    let mut mem = HashMap::new();
    let mut mask = Command::Mask {
        and: std::u64::MAX,
        or: 0u64,
        raw: vec![],
    };
    for inst in input {
        match inst {
            Command::Write { val, addr } => {
                mem.insert(*addr, mask.apply_mask_v1(*val));
            }
            m => {
                mask = m.clone();
            }
        }
    }
    return mem.values().copied().sum();
}

#[aoc(day14, part2)]
pub fn solve_part2(input: &[Command]) -> u64 {
    let mut mem = HashMap::new();
    let mut mask = Command::Mask {
        and: std::u64::MAX,
        or: 0u64,
        raw: vec![],
    };
    for inst in input {
        match inst {
            Command::Write { val, addr } => {
                for addr in mask.apply_mask_v2(*addr) {
                    mem.insert(addr, val);
                }
            }
            m => {
                mask = m.clone();
            }
        }
    }
    return mem.values().copied().sum();
}