zerosleeps

def get_raw_input()
  IO.read(File.join(__dir__, '../day_14_input.txt'))
end

def parse_raw_input(raw_input)
    raw_input.strip.lines(chomp: true)
end

def apply_bitmask(input, mask)
  ( input | mask.tr('X', '0').to_i(2) ) & mask.tr('X', '1').to_i(2)
end

def parse_instruction(instruction)
  instruction.match(/^mem\[(?<location>\d+)\] = (?<value>\d+)$/).named_captures
end

def apply_bitmask_v2(input, mask)
  input = input.to_s(2).rjust(36, '0')
  mask.chars.each_with_index do |c,i|
    case c
    when '1'
      input[i] = '1'
    when 'X'
      input[i] = 'X'
    end
  end

  permutations = [input]
  while permutations.any? { |e| e.include?('X') }
    perm = permutations.shift
    permutations << perm.sub(/X/, '0')
    permutations << perm.sub(/X/, '1')
  end
  permutations
end

def run(parsed_input, version=1)
  memory = {}
  mask = ''
  parsed_input.each do |line|
    if line =~ /^mask = /
      mask = line.split.last
    else
      instruction = parse_instruction(line)
      if version == 1
        memory[instruction['location']] = apply_bitmask(instruction['value'].to_i, mask)
      else
        apply_bitmask_v2(instruction['location'].to_i, mask).each do |location|
          memory[location.to_i(2)] = instruction['value'].to_i
        end
      end
    end
  end
  memory.values.sum
end

puts "Part one: #{run(parse_raw_input(get_raw_input))}"
puts "Part two: #{run(parse_raw_input(get_raw_input), 2)}"