use std::collections::HashMap;
use std::fs;

use itertools::Itertools;
use lazy_static::lazy_static;
use regex::Regex;

const MASK_36_BITS: u64 = 0xFFFFFFFFF;

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let input = fs::read_to_string("input")?;
    let program = parse_program(&input)?;

    // Part 1
    let memory = run_program_chip1(&program);
    println!("{}", memory.iter().map(|(_, value)| value).sum::<u64>());

    // Part 2
    let memory = run_program_chip2(&program);
    println!("{}", memory.iter().map(|(_, value)| value).sum::<u64>());

    Ok(())
}

#[derive(Copy, Clone, Debug)]
enum Instruction {
    Mask(u64, u64, u64),
    Mem(u64, u64),
}

impl Instruction {
    fn new(instruction: &str) -> Result<Self, String> {
        lazy_static! {
            static ref INSTRUCTION_RE: Regex =
                Regex::new(r"([[:alpha:]]+?)\[?(\d+?)?\]? = ([0-9X]+)")
                    .expect("Regex should compile");
        }

        let caps = INSTRUCTION_RE
            .captures(instruction)
            .ok_or("Instruction did not match")?;

        Ok(match &caps[1] {
            "mask" => {
                let and = u64::from_str_radix(&caps[3].replace("X", "1"), 2)
                    .map_err(|e| format!("Could not parse and: {}", e))?;
                let or = u64::from_str_radix(&caps[3].replace("X", "0"), 2)
                    .map_err(|e| format!("Could not parse or: {}", e))?;

                Instruction::Mask(and, or, and ^ or)
            }
            "mem" => {
                let location = caps[2]
                    .parse()
                    .map_err(|e| format!("Could not parse location: {}", e))?;
                let value = caps[3]
                    .parse()
                    .map_err(|e| format!("Could not parse value: {}", e))?;

                Instruction::Mem(location, value)
            }
            _ => Err("Invalid instruction")?,
        })
    }
}

fn parse_program(input: &str) -> Result<Vec<Instruction>, String> {
    input.lines().map(Instruction::new).collect()
}

fn run_program_chip1(program: &Vec<Instruction>) -> Vec<(u64, u64)> {
    let mut mask = (u64::max_value() & MASK_36_BITS, 0);

    program
        .iter()
        .filter_map(|instruction| match instruction {
            Instruction::Mask(and, or, _) => {
                mask = (*and, *or);
                None
            }
            Instruction::Mem(address, value) => Some((*address as u64, value & mask.0 | mask.1)),
        })
        .collect::<Vec<(u64, u64)>>()
        .into_iter()
        .rev()
        .unique_by(|address| address.0)
        .collect()
}

fn run_program_chip2(program: &Vec<Instruction>) -> Vec<(u64, u64)> {
    let mut memory = HashMap::new();
    let mut mask = (0, 0);

    for instruction in program {
        match instruction {
            Instruction::Mask(_, or, floating) => {
                mask = (*or, *floating);
            }
            Instruction::Mem(address, value) => {
                // We set floating bits to 0 here, so that they can be
                // toggled with addition
                let base = (address | mask.0) & !mask.1 & MASK_36_BITS;

                let mut floating = !mask.1 & MASK_36_BITS;
                while floating <= MASK_36_BITS {
                    // For each unset bit in the floating mask, we
                    // toggle the bit in the address
                    memory.insert(base | (!floating & MASK_36_BITS), *value);
                    floating += 1;

                    // Since add does overflows, we reset 1s to 0s if
                    // they're unset in the original mask
                    floating |= !mask.1 & MASK_36_BITS;
                }
            }
        }
    }

    memory.drain().collect()
}

#[cfg(test)]
mod tests {
    use super::*;
    use indoc::indoc;

    #[test]
    fn test_simple() -> Result<(), Box<dyn std::error::Error>> {
        let input = indoc!(
            "
             mask = XXXXXXXXXXXXXXXXXXXXXXXXXXXXX1XXXX0X
             mem[8] = 11
             mem[7] = 101
             mem[8] = 0
            "
        );

        let program = parse_program(input)?;
        let memory = run_program_chip1(&program);

        assert_eq!(memory.iter().map(|(_, value)| value).sum::<u64>(), 165);
        Ok(())
    }

    #[test]
    fn test_simple2() -> Result<(), Box<dyn std::error::Error>> {
        let input = indoc!(
            "
             mask = 000000000000000000000000000000X1001X
             mem[42] = 100
             mask = 00000000000000000000000000000000X0XX
             mem[26] = 1
            "
        );

        let program = parse_program(input)?;
        let memory = dbg!(run_program_chip2(&program));

        assert_eq!(memory.iter().map(|(_, value)| value).sum::<u64>(), 208);
        Ok(())
    }
}