#!/usr/bin/python3

import argparse
import glob
import json
import logging
import os
import shutil
import subprocess
import sys
import tempfile
import time
import yaml


OK_MESSAGE = """
Installation was successful. You can boot into the installed system now.
Remove the installer USB drive after the system has booted.
"""


verbose = False


def indent(s):
    return "".join(f"    {line}\n" for line in s.splitlines())


def log(msg):
    if verbose:
        print("INSTALLER:", msg)
    logging.info(msg)


def error(msg):
    if verbose:
        print("INSTALLER ERROR:", msg)
    logging.error(msg)


def run(argv, **kwargs):
    log(f"RUN: {argv[0]}")
    for arg in argv:
        log(f"RUN: - {arg!r}")
    for key in kwargs:
        log(f"RUN: - {key}={kwargs[key]!r}")
    check = False
    if kwargs.get("check"):
        check = True
        del kwargs["check"]
    if "capture_output" not in kwargs:
        kwargs["capture_output"] = True
    p = subprocess.run(argv, **kwargs)
    if check and p.returncode != 0:
        stdout = indent(p.stdout.decode())
        stderr = indent(p.stderr.decode())
        error(f"Ran command that failed: {argv}")
        error(f"Exit code {p.returncode}")
        error(f"Stdout:\n{stdout}")
        error(f"Stderr:\n{stderr}")
        sys.exit("command failed")
    return p


def physical_volumes():
    log("list physical volumes")
    p = run(["pvdisplay", "-C", "--noheadings"], capture_output=True, check=True)
    lines = p.stdout.decode().splitlines()
    pvs = []
    for line in lines:
        words = line.split()
        pv = words[0]
        vg = words[1]
        pvs.append({"pv": pv, "vg": vg})
    return pvs


def logical_volumes():
    log("list logical volumes")
    p = run(["lvdisplay", "-C", "--noheadings"], capture_output=True, check=True)
    lines = p.stdout.decode().splitlines()
    lvs = []
    for line in lines:
        words = line.split()
        lv = words[0]
        vg = words[1]
        path = os.path.realpath(f"/dev/{vg}/{lv}")
        lvs.append({"lv": lv, "vg": vg, "path": path})
    return lvs


def volume_groups():
    log("list volume groups")
    p = run(["vgdisplay", "-C", "--noheadings"], capture_output=True, check=True)
    lines = p.stdout.decode().splitlines()
    return [line.split()[0] for line in lines if line.strip()]


def mount_points():
    log("find mount all points")
    mounts = {}
    for line in open("/proc/mounts").readlines():
        words = line.split()
        dev = os.path.realpath(words[0]) if words[0].startswith("/") else words[0]
        mounts[dev] = {
            "mount": words[1],
            "type": words[2],
        }
    return mounts


def find_mount_points(mounts, dev):
    log(f"find mount points using {dev}")
    res = []
    if dev in mounts:
        root = mounts[dev]["mount"]
        for m in mounts.values():
            if m["mount"].startswith(root):
                res.append(m)
    return res


def is_luks(path):
    log(f"is {path} a LUKS device?")
    p = run(["cryptsetup", "isLuks", path], check=False)
    return p.returncode == 0


def clean_up_disks(drives):
    log(f"clean system LVM2 and drives from everything: {drives}")

    mounts = mount_points()
    pvs = physical_volumes()
    lvs = logical_volumes()
    vgs = volume_groups()

    log(f"PVs: {pvs}")
    log(f"VGs: {vgs}")
    log(f"LVs: {lvs}")

    for lv in lvs:
        log(f"LV: {lv}")
        for m in find_mount_points(mounts, lv["path"]):
            log(f"unmount {m['mount']}")
            run(["umount", m["mount"]], capture_output=True)

    for vg in vgs:
        log(f"remove volume group {vg}")
        run(["vgremove", "--yes", vg], check=True, capture_output=True)

    for pv in pvs:
        log(f"remove physical volume {pv}")
        run(["pvremove", pv["pv"]], check=True, capture_output=True)
        if is_luks(pv["pv"]):
            run(["cryptsetup", "close", pv["pv"]], check=True, capture_output=True)

    for mapping in glob.glob("/dev/mapper/*"):
        if not mapping.endswith("/control"):
            log(f"open LUKS volume {mapping} (just in case it is one)")
            run(["cryptsetup", "close", mapping], check=False, capture_output=True)

    for drive in drives:
        log(f"blkdiscard {drive}")
        p = run(
            ["blkdiscard", "--force", "--secure", drive],
            capture_output=True,
        )
        if p.returncode != 0:
            log("secure blkdiscard failed, trying insecure")
            run(
                ["blkdiscard", "--force", drive],
                check=True,
                capture_output=True,
            )


def mklabel(device):
    return {
        # Create a GPT partition table. It works well for all modern
        # systems.
        "mklabel": "gpt",
        "device": device,
    }


def mkpart(device, tag, start, end):
    return {
        "mkpart": "primary",
        "device": device,
        "start": start,
        "end": end,
        "tag": tag,
    }


def mkfs(tag, fstype):
    return {
        "mkfs": fstype,
        "partition": tag,
    }


def cryptsetup(tag, password, name):
    return {
        "cryptsetup": "cryptsetup0",
        "password": password,
        "name": name,
    }


def vgcreate(tag, drives):
    return {
        "vgcreate": tag,
        "physical": drives,
    }


def lvcreate(vg, name, size):
    return {
        "lvcreate": vg,
        "name": name,
        "size": size,
    }


def mount(tag, dirname=None, mount_on=None):
    d = {"mount": tag, "zerofree": False}
    if dirname is not None:
        d["dirname"] = dirname
    if mount_on is not None:
        d["mount-on"] = mount_on
    return d


def unpack_rootfs(tag):
    return {
        "unpack-rootfs": tag,
    }


def cache_rootfs(tag):
    return {
        "cache-rootfs": tag,
        "unless": "rootfs_unpacked",
    }


def debootstrap(debian_release, tag):
    return {
        "debootstrap": debian_release,
        "mirror": "http://deb.debian.org/debian",
        "target": tag,
        "unless": "rootfs_unpacked",
        "require_empty_target": False,
    }


def apt(tag, packages, unless=True):
    d = {
        "apt": "install",
        "packages": packages,
        "tag": tag,
    }
    if unless:
        d["unless"] = "rootfs_unpacked"
    return d


def virtual_filesystems(tag):
    return {
        "virtual-filesystems": "root",
    }


def fstab(tag):
    return {
        "fstab": tag,
    }


def grub(device, root, efi):
    return {
        "grub": "uefi",
        "tag": root,
        "efi": efi,
        "quiet": True,
        "image-dev": device,
    }


def vmdb_spec(system, ansible_vars, std_yml):
    device = "{{ image }}"
    steps = [
        # Create partition table.
        mklabel(device),
        # Create a partition for UEFI to boot from.
        mkpart(device, "efi", "0%", "500M"),
        # Create a separate /boot partition, in case we want to have LUKS.
        mkpart(device, "boot", "500M", "1G"),
        # Format the EFI partition. This MUST be vfat.
        mkfs("efi", "vfat"),
        # Format /boot. This is conventionally ext2. This file system
        # gets very little I/O, but it MUST be supported by GRUB, so
        # ext2 seems like a nice, safe choice.
        mkfs("boot", "ext2"),
    ]

    # Set up pv0 for lvm2, either encrypted or cleartext.
    if system.luks:
        steps.extend(
            [
                mkpart(device, "cryptsetup0", "1G", "100%"),
                cryptsetup("cryptsetup0", system.luks, "pv0"),
            ]
        )
        for i, drive in enumerate(system.extra_drives):
            steps.append(cryptsetup(f"cryptsetuo{i+1}", system.luks, f"pv{i+1}"))
    else:
        steps.append(mkpart(device, "pv0", "1G", "100%"))
        for i, drive in enumerate(system.extra_drives):
            steps.extend(
                [
                    mklabel(drive),
                    mkpart(drive, f"pv{i+1}", "0%", "100%"),
                ]
            )

    # Create file systems and install Debian.
    steps.extend(
        [
            # Create an LVM2 volume group using pv0, which we create
            # earlier. At this point, if LUKS is used, pv0 is the unlocked,
            # open, cleartext block device.
            vgcreate(
                "vg0",
                [f"pv{i}" for i in range(len([system.drive] + system.extra_drives))],
            ),
            # Create a 20 gigabyte LV for the root file system. That's big
            # enough for a desktop system.
            lvcreate("vg0", "root", "20G"),
            # format the root file system.
            mkfs("root", system.root_fstype),
            # Mount the root file system.
            mount("root"),
            # Mount /boot on top of the root file system.
            mount("boot", dirname="/boot", mount_on="root"),
            # Mount /boot/efi.
            mount("efi", dirname="/boot/efi", mount_on="boot"),
        ]
    )

    # Add any additional LVs.
    for lv in system.extra_lvs:
        steps.extend(
            [
                lvcreate("vg0", lv["name"], lv["size"]),
                mkfs(lv["name"], lv.get("fstype", "ext4")),
                mount(lv["name"], dirname=lv["mounted"], mount_on="root"),
            ]
        )

    steps.extend(
        [
            # If we have a cached version of the installed system, unpack
            # it now. Otherwise do nothing. Note that if you make any
            # changes to the steps marked "unless: rootfs_unpacked", you
            # have to remember to manually remove the cache file. v-i or
            # vmdb2 won't do that automatically for you.
            unpack_rootfs("root"),
            debootstrap(system.debian_release, "root"),
            apt(
                "root",
                [
                    "console-setup",
                    "dosfstools",
                    "linux-image-amd64",
                    "locales-all",
                    "lvm2",
                    "psmisc",
                    "python3",
                    "python3-apt",
                    "ssh",
                    "strace",
                ],
            ),
            # If we didn't unpack an existing cache archive, make one now.
            # Otherwise, skip this step.
            cache_rootfs("root"),
            # This MUST be after the debootstrap step.
            virtual_filesystems("root"),
            # Create /etc/fstab (and, if LUKS is used, /etc/crypttab).
            fstab("root"),
            # These MUST come after the fstab step so that they add the
            # crypttab in the initramfs. We install them regardless of
            # whether LUKS is used: they're harmless if LUKS isn't used.
            apt(
                "root",
                [
                    "cryptsetup",
                    "cryptsetup-initramfs",
                ],
                unless=False,
            ),
            # This also MUST come outside the rootfs caching, as it install
            # things outside the file systems, and those won't be in the
            # cache.
            grub(device, "root", "efi"),
        ]
    )

    # If playbooks have been specified, add ansible steps.
    for p in [std_yml] + system.extra_playbooks:
        if p:
            steps.append({"ansible": "root", "playbook": p, "extra_vars": ansible_vars})

    return {"steps": steps}


class Timings:
    def __init__(self):
        self.started = time.time()
        self.times = {}
        self.whats = []

    def reached(self, what):
        assert what not in self.whats
        assert what not in self.times
        self.times[what] = time.time()
        self.whats.append(what)

    def report(self):
        prev = self.started
        for what in self.whats:
            secs = "%.1f" % (self.times[what] - prev)
            prev = self.times[what]
            log(f"Duration: {what}: {secs} s")
            print(f"Duration: {what}: {secs} s")


def cache_name(spec):
    return f"{spec.debian_release}.cache.tar.gz"


class SystemSpec:
    def __init__(self, filename):
        REQUIRED = "required"
        self._obj = {
            "hostname": REQUIRED,
            "drive": REQUIRED,
            "extra_drives": [],
            "extra_lvs": [],
            "extra_playbooks": [],
            "ansible_vars": {},
            "ansible_vars_files": [],
            "luks": "",
            "debian_release": "stable",
            "root_fstype": "ext4",
        }
        with open(filename) as f:
            obj = yaml.safe_load(f)

        # Check for unknown keys.
        for key in obj:
            if key not in self._obj:
                sys.exit(f"spec has unknown key: {key}")

        # Check for missing required keys.
        for key in self._obj:
            if self._obj[key] == REQUIRED and key not in obj:
                sys.exit(f"spec lacks required key {key}")

        # Check for types of values.
        for key in self._obj:
            if key in obj:
                e = type(self._obj[key])
                a = type(obj[key])
                if a != e:
                    sys.exit(f"spec key {key} has unexpected type {a}, wanted {e}")

        self._obj.update(obj)

        for key in self._obj:
            setattr(self, key, self._obj[key])

        self._check()

    def __repr__(self):
        return repr(self.as_dict())

    def as_dict(self):
        return {key: self._obj[key] for key in self._obj}

    def as_json(self):
        return json.dumps(self.as_dict(), indent=4)

    def _check(self):
        try:
            for filename in self.extra_playbooks + self.ansible_vars_files:
                if not os.path.exists(filename):
                    raise Exception(
                        f"spec refers to file {filename} which does not exist"
                    )
        except Exception as e:
            error(e)
            sys.exit("problem with spec file")


def main():
    timings = Timings()

    p = argparse.ArgumentParser()
    p.add_argument(
        "-v", "--verbose", action="store_true", help="report to stdout what's happening"
    )
    p.add_argument(
        "--very-verbose", action="store_true", help="be more verbose to stdout"
    )
    p.add_argument(
        "--log", default="install.log", help="log to file LOG (default is no log file)"
    )
    p.add_argument(
        "--cache",
        default=None,
        help="use CACHE to cache debootstrap output (default is computed from spec file content)",
    )
    p.add_argument(
        "--std-yml",
        default="/usr/share/v-i/std.yml",
        help="Ansible playbook for standard system (%(default)s)",
    )
    p.add_argument("spec")
    args = p.parse_args()

    logging.basicConfig(
        filename=args.log,
        level=logging.DEBUG,
        format="%(asctime)s %(levelname)s %(message)s",
    )

    global verbose
    verbose = args.verbose

    log("v-i starts")

    system = SystemSpec(args.spec)
    spec_json = system.as_json()
    log(f"formatted spec:\n{spec_json}")

    cache = args.cache or cache_name(system)
    log(f"cache: {cache}")

    ansible_vars = dict(system.ansible_vars)
    ansible_vars["hostname"] = system.hostname
    for filename in system.ansible_vars_files:
        log(f"reading Ansible vars from {filename}")
        with open(filename) as f:
            vars_dict = yaml.safe_load(f)
        ansible_vars.update(vars_dict)
    ansible_vars["debian_release"] = system.debian_release
    ansible_vars_json = json.dumps(ansible_vars, indent=4)
    log(f"ansible_vars:\n{ansible_vars_json}")
    timings.reached("read configuration")

    clean_up_disks([system.drive] + system.extra_drives)
    timings.reached("clean up storage")

    vmdb = vmdb_spec(system, ansible_vars, args.std_yml)
    x = yaml.safe_dump(vmdb, indent=4)
    log(f"vmdb spec:\n{x}")
    tmp = tempfile.mkdtemp()
    specfile = os.path.join(tmp, "spec.yaml")
    if args.very_verbose:
        yaml.dump(vmdb, stream=sys.stdout, indent=4)
    with open(specfile, "w") as f:
        yaml.dump(vmdb, stream=f, indent=4)

    log(f"run vmdb2 to install on {system.drive}")
    env = dict(os.environ)
    env["ANSIBLE_STDOUT_CALLBACK"] = "yaml"
    env["ANSIBLE_NOCOWS"] = "1"
    env["ANSIBLE_LOG_PATH"] = "ansible.log"

    argv = [
        "vmdb2",
        f"--rootfs-tarball={cache}",
        f"--log={args.log}",
        f"--image={system.drive}",
        specfile,
    ]
    if verbose:
        argv.append("--verbose")
    run(argv, check=True, capture_output=True, env=env)
    timings.reached("vmdb2")

    log("cleanup")
    shutil.rmtree(tmp)
    timings.reached("clean up")

    log("OK, done")
    print("OK, done")

    timings.report()

    print(OK_MESSAGE)


main()