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

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, :device, :fs_uuid
  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
    @device = nil
    @fs_uuid = nil
  end
end

class ImageCreator
  attr_accessor :name, :dirs, :src_host, :img_path_base
  MiB = 1024 ** 2
  GiB = 1024 ** 3

  def initialize(name, dirs, src_host: nil)
    @name = name
    @dirs = dirs
    @src_host = src_host || name
    @img_path_base = "#{name}_#{Time.now.strftime '%FT%T%z'}"
  end

  def imgpath(idx)
    "#{img_path_base}_#{idx}.img"
  end

  def create_disk
    dirs.each_with_index do |di, idx|
      _create_disk imgpath(idx), di, idx != 0
    end
  end

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

  def create_fs
    with_loopdev do |devices|
      dirs.each_with_index do |di, index|
        dev = di.device
        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}"
        system "e2label", dev, di.path == '/' ? 'ROOT' : di.path[1..-1].tr('/', '-')
        di.fs_uuid = `blkid -o value -s UUID #{di.device}`.chomp("\n")
      end
    end
  end

  def sync_dirs
    with_loopdev do |devices|
      devices.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}..."
          unless system("rsync", "-azHSAX", "--numeric-ids", "--info=progress2", "#{src_host}:#{di.srcpath}/", "#{mount_point}/", *((["--exclude"] * di.exclude.size).zip(di.exclude).flatten))
            warn "rsync exit with error, file transfer may not be completed."
          end
        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|
      with_loopdev do |devices|
        puts "Override grub with host version..."
        root_dev = "/dev/#{devices.first[/loop\d+/]}"
        rootfs_uuid = dirs.find { |d| d.path == '/'}.fs_uuid
        puts "New rootfs UUID=#{rootfs_uuid}"
        begin
          system("mount", devices.first, dir) or raise "Failed to mount #{devices.first} to #{dir}"
          system("mount", "--bind", "/dev", "#{dir}/dev") or raise "Failed to mount /dev to #{dir}/dev"
          system("mount", "--bind", "/proc", "#{dir}/proc") or raise "Failed to mount /proc to #{dir}/proc"

          dirs[1..-1].each_with_index do |di, idx|
            system "mkdir", "-p", "#{dir}#{di.path}"
            system("mount", di.device, "#{dir}#{di.path}") or raise "Failed to mount #{di.device} to #{dir}#{path}"
          end

          system "rm", "-f", "#{dir}/etc/systemd/system/udev.service", "#{dir}/etc/systemd/system/systemd-udevd.service"

          puts "Rewrite fstab..."
          File.open "#{dir}/etc/fstab", "w" do |f|
            dirs.each_with_index do |di, idx|
              f << %W(UUID=#{di.fs_uuid} #{di.path} ext4 defaults,noatime 0 #{di.path == '/' ? 1 : 2}).join("\t")
              f << "\n"
            end
          end

          system("chroot", dir, "apt-get", "-qy", "update")
          if File.read("#{dir}/etc/debian_version").to_f >= 9.0
            # Note: 2019-10-08 時点で Debian9 のカーネルバージョンに対応していないので、エラー回避のために既存のカーネルを全て削除し、強制的に jessie のカーネルをイントールする
            system("rm", "-f", "#{dir}/var/lib/dpkg/info/linux-image-#{`uname -r`.chomp}.prerm")
            system({'DEBIAN_FRONTEND' => 'noninteractive'}, "chroot", dir, "apt", "remove", "--purge", "-y", "linux-image-*")
            system("wget", "-O", "#{dir}/tmp/linux.deb", "http://security-cdn.debian.org/debian-security/pool/updates/main/l/linux/linux-image-3.16.0-10-amd64_3.16.81-1_amd64.deb") or raise "Failed to get jessie kernel"
            system("chroot", dir, "dpkg", "-i", "/tmp/linux.deb")
            system({'DEBIAN_FRONTEND' => 'noninteractive'}, "chroot", dir, "apt", "install", "-f", "-y") or raise "Failed to install jessie kernel"
          else
            system({'DEBIAN_FRONTEND' => 'noninteractive'}, "chroot", dir, "apt-get", "-y", "install", "linux-image-amd64")
          end

          puts "Update grub..."
          if File.exists? "#{dir}/var/lib/dpkg/info/grub-efi-amd64.list"
            system({'DEBIAN_FRONTEND' => 'noninteractive'}, "chroot", dir, "apt", "remove", "--purge", "-y", "grub-efi-amd64") or raise "Failed to purge grub-efi"
          end
          unless File.exists? "#{dir}/var/lib/dpkg/info/grub-pc.list"
            system("chroot", dir, "apt-get", "-qy", "update") or raise "Failed to install grub-pc"
            system({'DEBIAN_FRONTEND' => 'noninteractive'}, "chroot", dir, "apt-get", "-y", "install", "grub-pc")
          end
          File.open "#{dir}/boot/grub/device.map", "w" do |f|
            f.puts "(hd0)\t#{root_dev}"
          end
          system("chroot", dir, "grub-mkconfig", "-o", "/boot/grub/grub.cfg") or raise "grub-mkconfig fails."
          system(*%W(grub-install --no-floppy --grub-mkdevicemap=#{dir}/boot/grub/device.map --root-directory=#{dir} #{root_dev})) or raise "grub-install failed."
        ensure
          system("umount", "#{dir}/dev")
          system("umount", "#{dir}/proc")
          dirs.reverse[0..-2].each do |di, idx|
            system("umount", "#{dir}#{di.path}")
          end
          system("umount", dir)
        end
      end
    end
  end

  def write_json
    jdef = []
    dirs.each_with_index do |dir, idx|
      jdef.push({
        "Description" => dir.path == '/' ? 'root' : dir.path[1..-1].tr('/', '-'),
        "Format" => "raw",
        "UserBucket" => {
          "S3Bucket" => ENV.fetch("S3_BUCKET", "osdn-base-images"),
          "S3Key" => "#{ENV.fetch("S3_KEY_PREFIX", "src-disks/")}#{img_path_base}_#{idx}.img"
        }
      })
    end
    File.write "#{img_path_base}.json", JSON.pretty_generate(jdef)
  end


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

if $0 == __FILE__
  require 'optparse'

  opts = ARGV.getopts('l:', 'limit:', 'skip:')
  limit_pat = opts['limit'] || opts['l']
  limit_pat and
    limit_pat = Regexp.new(limit_pat)

  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|
        opts.kind_of?(Hash) or opts = {size: opts}
        dirs << SyncDirDef.new({path: path}.merge(opts.keys.map(&:to_sym).zip(opts.values).to_h))
      end
    else
      name = imgdef
    end
    if limit_pat
      limit_pat.match?(name) or next
    end
    dirs.empty? and dirs << SyncDirDef.new
    ImageCreator.new(name, dirs).run
  end
end