use std::collections::HashMap;
use std::error::Error;
use std::fs::File;
use std::io::{self, BufRead};
use std::vec::Vec;

#[macro_use]
extern crate lazy_static;
extern crate regex;

use regex::Regex;

enum Command {
    Mem(u64, u64),
    Mask(String),
}

impl Command {
    fn from_string(line: &str) -> Command {
        lazy_static! {
            static ref MASK_RE: Regex = Regex::new(r"^mask = ([X01]{36})$").unwrap();
            static ref MEM_RE: Regex = Regex::new(r"^mem\[([0-9]*)\] = ([0-9]*)$").unwrap();
        }

        if let Some(cap) = MEM_RE.captures_iter(line).nth(0) {
            let addr: u64 = cap[1].parse().unwrap();
            let value: u64 = cap[2].parse().unwrap();
            Command::Mem(addr, value)
        } else if let Some(cap) = MASK_RE.captures_iter(line).nth(0) {
            let mask = cap[1].to_string();
            Command::Mask(mask)
        } else {
            panic!("Unable to parse line: {}", line);
        }
    }
}

struct VM {
    mask: String,
    memory: HashMap<u64, u64>,
}

impl VM {
    fn new() -> VM {
        VM {
            mask: "000000000000000000000000000000000000".to_string(),
            memory: HashMap::new(),
        }
    }

    fn compute_addresses(&self, addr: u64) -> Vec<u64> {
        fn enumerate_addresses(mask: String, bit: usize, addr: u64) -> Vec<u64> {
            if bit == 0 {
                match mask.chars().nth(0).unwrap() {
                    '0' => {
                        return vec![addr & 1];
                    }
                    '1' => {
                        return vec![1];
                    }
                    'X' => {
                        return vec![0, 1];
                    }
                    c => panic!("Unsupported mask bit: {}", c),
                };
            } else {
                let lower_bits = enumerate_addresses(mask.clone(), bit - 1, addr);
                let addr_bit = addr & (1 << bit as u64);
                match mask.chars().nth(bit).unwrap() {
                    '0' => return lower_bits.iter().map(|b| b | addr_bit).collect(),
                    '1' => return lower_bits.iter().map(|b| b | (1 << bit as u64)).collect(),
                    'X' => {
                        let mut bit_zero = lower_bits.clone();
                        let mut bit_one: Vec<u64> =
                            lower_bits.iter().map(|b| b | (1 << bit as u64)).collect();
                        bit_zero.append(&mut bit_one);
                        return bit_zero;
                    }
                    c => panic!("Unsupported mask bit: {}", c),
                };
            }
        }

        enumerate_addresses(self.mask.chars().rev().collect(), 35, addr)
    }

    fn run(&mut self, programm: &Vec<Command>) {
        for inst in programm {
            match inst {
                Command::Mask(mask) => {
                    self.mask = mask.clone();
                }
                Command::Mem(addr, value) => {
                    let addrs = self.compute_addresses(*addr);
                    for address in addrs {
                        self.memory.insert(address, *value);
                    }
                }
            }
        }
    }

    fn compute_sum(&self) -> u64 {
        self.memory.values().fold(0, |acc, x| acc + x)
    }
}

fn main() -> Result<(), Box<dyn Error>> {
    let file = File::open("inputs/day14.txt")?;
    let lines = io::BufReader::new(file).lines().map(|l| l.unwrap());

    let programm: Vec<Command> = lines.map(|l| Command::from_string(&l)).collect();

    let mut vm = VM::new();
    vm.run(&programm);

    println!("Answer: {}", vm.compute_sum());

    Ok(())
}