#!/usr/bin/ruby
require 'pp'
require 'shellwords'
require 'tmpdir'
require 'yaml'

class SyncDirDef
  DEFAULT_EXCLUDE = %w[/proc/* /sys/* /dev/mqueue /dev/hugepages /run/* /var/lib/os-prober/mount /swap /dev/shm/* /var/lib/lxcfs/*]
  attr_accessor :path, :size, :exclude, :size, :srcpath, :fs_features
  def initialize(path: '/', size: 8, exclude: DEFAULT_EXCLUDE, srcpath: nil, fs_features: nil)
    @path = path
    @size = size.to_f
    @exclude = (DEFAULT_EXCLUDE + [*exclude]).uniq
    @srcpath = srcpath || path
    @fs_features = fs_features
  end
end


class ImageCreator
  attr_accessor :name, :dirs, :src_host

  def initialize(name, dirs, src_host: nil)
    @name = name
    @dirs = dirs
    @src_host = src_host || name
  end

  def size
    @size and return @size 
    @size = 3 * (1024**2) # alignment + BIOS boot partition + gap
    dirs.each do |di|
      @size += di.size * (1024**3)
    end
    @size
  end

  def imgpath
    @imgpath ||= "#{name}_#{Time.now.strftime '%FT%T%z'}.img"
  end

  def create_disk
    raise "Disk image #{imgpath} is already exists!" if File.exists? imgpath
    puts "Creating disk image #{imgpath} (#{'%5.2f' % (size.to_f/(1024**3))} GiB)..."
    last_mb = 1
    File.open(imgpath, "w") do |f|
      f.truncate size
    end
    system("parted", "-s", imgpath, "mklabel", "gpt") or raise "Failed to create partition label"
    system("parted", "-s", imgpath, "mkpart", "primary", "#{last_mb}MiB", "#{last_mb+1}MiB") or raise "Failed to create boot partition"
    system("parted", "-s", imgpath, "set", "1", "bios_grub", "on") or raise "Failed to set bios boot partition"
    last_mb += 1
    dirs.each do |di|
      system("parted", "-s", imgpath, "mkpart", "primary", "#{last_mb}MiB", "#{di.size * 1024 + last_mb}MiB") or raise "Failed to create partition: #{l}"
      last_mb += di.size * 1024
    end
    puts "Image partition created."
    #system "sfdisk", "-l", imgpath
  end
  
  def with_loopdev &block
    begin
      system("kpartx", "-as", imgpath) or raise "Failed to map loop device"
      maplines = `kpartx -l #{Shellwords.escape imgpath}`.split("\n")
      devices = []
      maplines.each do |map|
        devices << "/dev/mapper/#{map[/loop\d+p\d+/]}"
      end
      yield devices, maplines.first[%r{/dev/loop\d+}]
    ensure
      system "kpartx", "-d", imgpath, err: "/dev/null"
    end
  end

  def create_fs
    with_loopdev do |devices, root|
      devices[1..-1].each_with_index do |(dev, _), index|
        puts "Creating filesystem on #{dev}..."
        cmd = %w(mkfs.ext4 -q)
        di = dirs[index]
        if di.fs_features
          cmd << '-O' << di.fs_features
        end
        cmd << dev
        system(*cmd) or raise "Failed to create file system on #{dev}"
      end
    end
  end

  def sync_dirs
    with_loopdev do |devices, root|
      devices[1..-1].each_with_index do |dev, idx|
        di = dirs[idx]
        mount_point = "/mnt/ci-#{$$}-#{name}-#{idx}"
        system("mkdir", "-p", mount_point)
        begin
          system("mount", dev, mount_point) or raise "Failed to mount file system #{dev} on #{mount_point}"
          puts "Copying #{src_host}:#{di.srcpath} to #{dev}..."
          system("rsync", "-azHSAX", "--numeric-ids", "--info=progress2", "#{src_host}:#{di.srcpath}/", "#{mount_point}/", *((["--exclude"] * di.exclude.size).zip(di.exclude).flatten)) or raise "rsync fails"
        ensure
          system("umount", mount_point)
          File.directory?(mount_point) and
            Dir.rmdir mount_point
        end
      end
    end
  end

  def fix_boot
    puts "Fixing boot environments..."
    Dir.mktmpdir("ci-#{$$}-#{name}") do |dir|
      system("cp", "-a", "/usr/lib/grub/i386-pc/boot.img", dir) or raise "Failed to copy boot.img"
      coreimg = "#{dir}/core.img"
      system("grub-mkimage", "-o", coreimg, "-O", "i386-pc", "-p", "(hd0,gpt2)/boot/grub", "biosdisk", "part_gpt", "ext2", "gzio", "xzio", "lzopio") or raise "Failed to create grub core image."
      with_loopdev do |devices, root|
        puts "Override grub with host version..."
        system("grub-bios-setup", "-d", dir, root) or raise "Failed to run grub-bios-setup"
        rootfs_uuid=`blkid -o value -s UUID #{devices[1]}`.chomp("\n")
        puts "New rootfs UUID=#{rootfs_uuid}"
        begin
          system("mount", devices[1], dir)

          puts "Rewrite fstab..."
          fstab = File.read "#{dir}/etc/fstab"
          fstab.gsub!(%r{^(UUID=|/)\S+(\s+/\s+)}, "UUID=#{rootfs_uuid}\\2")
          fstab.gsub!(%r{^(\S+\s+\S+\s+\S+\s+sw(?=\b))}, '#\1')
          File.write "#{dir}/etc/fstab", fstab

          unless File.exists? "#{dir}/vmlinuz"
            system("chroot", dir, "apt-get", "-qy", "update")
            system("chroot", dir, "apt-get", "-y", "install", "linux-image-amd64")
          end

          puts "Update grub.cfg..."
          system("mkdir", "-p", "#{dir}/boot/grub/i386-pc") or raise "Failed to create grub dir"
          system("cp -a /usr/lib/grub/i386-pc/*.mod #{dir}/boot/grub/i386-pc/") or raise "Failed to copy grub modules"
          if File.exists? "#{dir}/boot/grub/grub.cfg"
            grubconf = File.read "#{dir}/boot/grub/grub.cfg"
            if old_uuid = grubconf[/root=UUID=(\S+)/, 1]
              File.write "#{dir}/boot/grub/grub.cfg", grubconf.gsub(/#{old_uuid}/, rootfs_uuid)
            end
          else
            File.write "#{dir}/boot/grub/grub.cfg", <<-EOC
              set timeout=5
              insmod part_gpt
              insmod ext2
              insmod linux
              search --no-floppy --fs-uuid --set=root #{rootfs_uuid}
              menuentry 'Linux' {
                linux   /vmlinuz root=UUID=#{rootfs_uuid} ro
                initrd  /initrd.img
              }
            EOC
          end
        ensure
          system("umount", dir)
        end
      end
    end
  end

  def run
    create_disk
    create_fs
    sync_dirs
    fix_boot
    puts "Image creation has been complated (#{name})"
  end

end

if $0 == __FILE__
  list = YAML.load_file(ARGV[0] || 'image-list.yml')
  list.each do |imgdef|
    name = nil
    dirs = []
    if imgdef.kind_of?(Hash)
      name = imgdef['name']
      (imgdef['dirs'] || {}).each do |path, opts|
        dirs << SyncDirDef.new({path: path}.merge(opts.keys.map(&:to_sym).zip(opts.values).to_h))
      end
    else
      name = imgdef
    end
    dirs.empty? and dirs << SyncDirDef.new
    ImageCreator.new(name, dirs).run
  end
end