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

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

            let mut set_mask: u64 = 0;
            let mut unset_mask: u64 = 0;
            let mut bit: u64 = 1 << 35;

            for c in mask.chars() {
                match c {
                    '0' => unset_mask |= bit,
                    '1' => set_mask |= bit,
                    'X' => {}
                    _ => panic!("Unsupported mask char: {}", c),
                }
                bit = bit >> 1;
            }

            Command::Mask(set_mask, unset_mask)
        } else {
            panic!("Unable to parse line: {}", line);
        }
    }
}

struct VM {
    set_mask: u64,
    unset_mask: u64,
    memory: HashMap<u64, u64>,
}

impl VM {
    fn new() -> VM {
        VM {
            set_mask: 0,
            unset_mask: 0,
            memory: HashMap::new(),
        }
    }

    fn run(&mut self, programm: &Vec<Command>) {
        for inst in programm {
            match inst {
                Command::Mask(set, unset) => {
                    self.set_mask = *set;
                    self.unset_mask = *unset;
                }
                Command::Mem(addr, value) => {
                    let set_value = value | self.set_mask;
                    let unset_value = set_value & !self.unset_mask;
                    self.memory.insert(*addr, unset_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(())
}