#!/usr/bin/python3
"""
Assemble tree into a raw filesystem image

Assemble the tree into a raw filesystem image named `filename`, with the UUID
`root_fs_uuid`.

The image is a sparse file of the given `size`, which is created using the
`truncate(1)` command. The `size` is an integer with an optional suffix:
K,M,G,T,... (for powers of 1024) or KB,MB,GB,TB,... (powers of 1000).

NOTE: If the tree contents are larger than `size`, this assembler will fail.
On the other hand, since the image is a sparse file, the unused parts of the
image take up almost no disk space - so a 1GB tree in a 20GB image should not
use much more than 1GB disk space.

The filesystem UUID should be a standard (RFC4122) UUID, which you can
generate with uuid.uuid4() in Python, `uuidgen(1)` in a shell script, or
read from `/proc/sys/kernel/random/uuid` if your kernel provides it.
"""


import contextlib
import os
import subprocess
import sys

import osbuild.api
import osbuild.remoteloop as remoteloop


SCHEMA = """
"additionalProperties": false,
"required": ["filename", "root_fs_uuid", "size"],
"properties": {
  "filename": {
    "description": "Raw filesystem image filename",
    "type": "string"
  },
  "root_fs_uuid": {
    "description": "UUID for the filesystem",
    "type": "string",
    "pattern": "^[0-9A-Za-z]{8}(-[0-9A-Za-z]{4}){3}-[0-9A-Za-z]{12}$",
    "examples": ["9c6ae55b-cf88-45b8-84e8-64990759f39d"]
  },
  "size": {
    "description": "Maximum size of the filesystem",
    "type": "integer"
  },
  "fs_type": {
    "description": "Filesystem type",
    "type": "string",
    "enum": ["ext4", "xfs", "btrfs"],
    "default": "ext4"
  }
}
"""

@contextlib.contextmanager
def mount(source, dest, *options):
    os.makedirs(dest, 0o755, True)
    subprocess.run(["mount", *options, source, dest], check=True)
    try:
        yield
    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_btrfs(device, uuid):
    subprocess.run(["mkfs.btrfs", "-U", uuid, device], encoding='utf-8', check=True)


def main(tree, output_dir, options, loop_client):
    filename = options["filename"]
    root_fs_uuid = options["root_fs_uuid"]
    size = options["size"]
    fs_type = options.get("fs_type", "ext4")

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

    subprocess.run(["truncate", "--size", str(size), image], check=True)

    if fs_type == "ext4":
        mkfs_ext4(image, root_fs_uuid)
    elif fs_type == "xfs":
        mkfs_xfs(image, root_fs_uuid)
    elif fs_type == "btrfs":
        mkfs_btrfs(image, root_fs_uuid)
    else:
        raise ValueError("`fs_type` must be ext4, xfs or btrfs")

    # Copy the tree into the target image
    with loop_client.device(image) as loop, mount(loop, mountpoint):
        subprocess.run(["cp", "-a", f"{tree}/.", mountpoint], check=True)

    subprocess.run(["mv", image, f"{output_dir}/{filename}"], check=True)


if __name__ == '__main__':
    args = osbuild.api.arguments()
    args_input = args["inputs"]["tree"]["path"]
    args_output = args["tree"]
    r = main(args_input, args_output, args["options"], remoteloop.LoopClient("/run/osbuild/api/remoteloop"))
    sys.exit(r)