#!/usr/bin/python3

import contextlib
import json
import os
import socket
import shutil
import subprocess
import sys
import tempfile
import osbuild.remoteloop as remoteloop

STAGE_DESC = "Assemble a bootable partitioned disk image with qemu-img"
STAGE_INFO = """
Assemble a bootable partitioned disk image using `qemu-img`.

Creates a sparse MBR-partitioned disk image of the given `size`, with a single
bootable partition containing the root filesystem.

Installs GRUB2 (using the buildhost's `/usr/lib/grub/i386-pc/boot.img` etc.) as
the bootloader.

Copies the tree contents into the root filesystem and then converts the raw
sparse image into the format requested with the `fmt` option.

Buildhost commands used: `truncate`, `mount`, `umount`, `sfdisk`,
`grub2-mkimage`, `mkfs.ext4` or `mkfs.xfs`, `qemu-img`.
"""
STAGE_OPTS = """
"required": ["format", "filename", "ptuuid", "root_fs_uuid", "size"],
"properties": {
  "format": {
    "description": "Image file format to use",
    "type": "string",
    "enum": ["raw", "qcow2", "vdi", "vmdk", "vpc"]
  },
  "filename": {
    "description": "Image filename",
    "type": "string"
  },
  "ptuuid": {
    "description": "UUID for the disk image's partition table",
    "type": "string"
  },
  "root_fs_uuid": {
    "description": "UUID for the root filesystem",
    "type": "string"
  },
  "size": {
    "description": "Virtual disk size",
    "type": "string"
  },
  "root_fs_type": {
    "description": "Type of the root filesystem",
    "type": "string",
    "enum": ["ext4", "xfs"],
    "default": "ext4"
  }
}
"""

@contextlib.contextmanager
def mount(source, dest):
    subprocess.run(["mount", source, dest], check=True)
    try:
        yield dest
    finally:
        subprocess.run(["umount", "-R", dest], check=True)


def mkfs_ext4(device, uuid):
    subprocess.run(["mkfs.ext4", "-U", uuid, device], input="y", encoding='utf-8', check=True)


def mkfs_xfs(device, uuid):
    subprocess.run(["mkfs.xfs", "-m", f"uuid={uuid}", device], encoding='utf-8', check=True)


def mkfs_for_type(device, uuid, fs_type):
    if fs_type == "ext4":
        maker = mkfs_ext4
    elif fs_type == "xfs":
        maker = mkfs_xfs
    else:
        raise ValueError("Unknown filesystem type")
    maker(device, uuid)


def create_partition_table(image, options):
    """Set up the partition table of the image"""
    ptuuid = options["ptuuid"]
    root_fs_uuid = options["root_fs_uuid"]
    root_fs_type = options.get("root_fs_type", "ext4")

    partition_table = f"label: mbr\nlabel-id: {ptuuid}\nbootable, type=83"
    subprocess.run(["sfdisk", "-q", image], input=partition_table, encoding='utf-8', check=True)

    r = subprocess.run(["sfdisk", "--json", image], stdout=subprocess.PIPE, encoding='utf-8', check=True)
    partition_table = json.loads(r.stdout)

    partition = partition_table["partitiontable"]["partitions"][0]
    partitions = [{
        "start": partition["start"] * 512,
        "size": partition["size"] * 512,
        "filesystem": {
            "type": root_fs_type,
            "uuid": root_fs_uuid,
            "mountpoint": "/"
        }
    }]

    return partitions


def install_grub2(image, fs_module, partition_offset):
    """Install grub2 to image"""
    grub2_core = "/var/tmp/grub2-core.img"

    # Create the level-2 bootloader
    # The purpose of this is to find the grub modules and configuration
    # to be able to start the level-3 bootloader. It contains the modules
    # necessary to do this, but nothing else.
    subprocess.run(["grub2-mkimage",
                    "--verbose",
                    "--directory", "/usr/lib/grub/i386-pc",
                    "--prefix", "(,msdos1)/boot/grub2",
                    "--format", "i386-pc",
                    "--compression", "auto",
                    "--output", grub2_core,
                    "part_msdos", fs_module, "biosdisk"],
                   check=True)

    assert os.path.getsize(grub2_core) < partition_offset - 512

    with open(image, "rb+") as image_f:
        # Install the level-1 bootloader into the start of the MBR
        # The purpose of this is simply to jump into the level-2 bootloader.
        with open("/usr/lib/grub/i386-pc/boot.img", "rb") as boot_f:
            # The boot.img file is 512 bytes, but we must only copy the first 440
            # bytes, as these contain the bootstrapping code. The rest of the
            # first sector contains the partition table, and must not be
            # overwritten.
            image_f.write(boot_f.read(440))

        # Install the level-2 bootloader into the space after the MBR, before
        # the first partition.
        with open(grub2_core, "rb") as core_f:
            image_f.seek(512)
            shutil.copyfileobj(core_f, image_f)


def main(tree, output_dir, options, loop_client):
    fmt = options["format"]
    filename = options["filename"]
    size = options["size"]
    root_fs_type = options.get("root_fs_type", "ext4")

    # sfdisk works on sectors of 512 bytes and ignores excess space - be explicit about this
    if size % 512 != 0:
        raise ValueError("`size` must be a multiple of sector size (512)")

    if fmt not in ["raw", "raw.xz", "qcow2", "vdi", "vmdk", "vpc"]:
        raise ValueError("`format` must be one of raw, qcow, vdi, vmdk, vpc")

    if root_fs_type == "ext4":
        mkfs = mkfs_ext4
        grub2_fs_module = "ext2"
    elif root_fs_type == "xfs":
        mkfs = mkfs_xfs
        grub2_fs_module = "xfs"
    else:
        raise ValueError("`root_fs_type` must be either ext4 or xfs")

    image = "/var/tmp/osbuild-image.raw"

    # Create an empty image file
    subprocess.run(["truncate", "--size", str(size), image], check=True)

    # The partition table
    partitions = create_partition_table(image, options)
    partition_offset = partitions[0]["start"]

    # Create the level-2 bootloader
    install_grub2(image, grub2_fs_module, partition_offset)

    # Now assemble the filesystem hierarchy and copy the tree into the image
    with contextlib.ExitStack() as cm:
        root = cm.enter_context(tempfile.TemporaryDirectory(prefix="osbuild-mnt"))
        # sort the partition according to their position in the filesystem tree
        for partition in sorted(partitions, key=lambda p: len(p["filesystem"]["mountpoint"])):
            offset, size = partition["start"], partition["size"]
            filesystem = partition["filesystem"]
            loop = cm.enter_context(loop_client.device(image, offset, size))
            # make the specified filesystem
            mkfs_for_type(loop, filesystem["uuid"], filesystem["type"])
            # now mount it
            mountpoint = os.path.normpath(f"{root}/{filesystem['mountpoint']}")
            os.makedirs(mountpoint, exist_ok=True)
            cm.enter_context(mount(loop, mountpoint))
        # the filesystem tree should now be properly setup,
        # copy the tree into the target image
        subprocess.run(["cp", "-a", f"{tree}/.", root], check=True)

    if fmt == "raw":
        subprocess.run(["cp", image, f"{output_dir}/{filename}"], check=True)
    elif fmt == "raw.xz":
        with open(f"{output_dir}/{filename}", "w") as f:
            subprocess.run(["xz", "--keep", "--stdout", "-0", image], stdout=f, check=True)
    else:
        extra_args = {
            "qcow2": ["-c"],
            "vdi": [],
            "vmdk": ["-c"],
            "vpc": ["-o", "subformat=fixed,force_size"]
        }
        subprocess.run([
            "qemu-img",
            "convert",
            "-O", fmt,
            *extra_args[fmt],
            image,
            f"{output_dir}/{filename}"
        ], check=True)


if __name__ == '__main__':
    args = json.load(sys.stdin)

    with socket.socket(socket.AF_UNIX, socket.SOCK_DGRAM) as sock:
        sock.setsockopt(socket.SOL_SOCKET, socket.SO_PASSCRED, 1)
        sock.connect("/run/osbuild/api/remoteloop")
        ret = main(args["tree"], args["output_dir"], args["options"], remoteloop.LoopClient(sock))

    sys.exit(ret)