require 'facter'

Facter.add(:wmf__disks) do
  confine :kernel => 'Linux'
  disk_dir = '/sys/block'
  disks = {ssd: [], hdd: [], virtual: []}
  virtual = Facter.fact('virtual').value
  setcode do
    if virtual == 'physical'
      Dir.each_child(disk_dir) do |disk|
        type = File.read(File.join(disk_dir, disk, 'queue/rotational')).strip
        if type == '0'
          disks[:ssd] << disk
        else
          disks[:hdd] << disk
        end
      end
      disks[:type] = disks[:hdd].length > disks[:ssd].length ? 'hdd' : 'ssd'
    else
      disks[:virtual] = Dir.children(disk_dir)
      disks[:type] = virtual
    end
    disks
  end
end