#!/usr/bin/env python3

# OS dependencies:
# - runs only on Linux
# - must be a user allowed to "sudo"
# - qemu-utils
# - lvm2
# - parted
# - The LVM volume group name "system" must be available

import os
import time
import subprocess
import getpass

import helpers


env = helpers.create_env()

base_path = env.base_path
build_path = env.build_path
mount_path = env.mount_path
config = env.disk_config

_run = helpers.run
_sudo = helpers.sudo
_chroot_cmds = helpers.chroot_cmds

image_filename = os.path.join(build_path, config.image_file)

mountpoints = [
    (volume['mountpoint'], '/dev/{}/{}'.format(
        config.volgroup, volume['name']
    ))
    for volume in config.volumes if volume['name'] != 'swap'
]
mountpoints.append(('/boot', '/dev/nbd0p1'))

chroot_mountpoints = [
    '/dev/', '/dev/pts', '/proc', '/sys',
]


def make_image():
    "Generate image file"
    print("Generate image file")
    subprocess.run([
        'qemu-img', 'create', '-f', 'qcow2', '-o', 'preallocation=metadata',
        image_filename, config.image_size,
    ], check=True)
    subprocess.run(['sync', ], check=True)


def make_filesystems():
    "Create partitions and filesystems"
    print("Partition disk")
    # TODO: Replace msdos image with gpt for very large disks
    labeltype = 'msdos'  # or 'gpt'
    # TODO: gpt doesn't provide enough room for GRUB! One additional
    # partition is required for this purpose.
    time.sleep(1)
    subprocess.run([
        'sudo', 'parted', '-s', '/dev/nbd0', 'mklabel', labeltype
    ], check=True)
    # For EFI: sudo parted -s /dev/nbd0 mkpart primary fat32 0% 200M
    # TODO: For hardware, the target disk must be replaced here
    target_disk = '/dev/nbd0'
    subprocess.run([
        'sudo', 'parted', '-s', target_disk,
        'mkpart', 'primary', 'ext4', '0%', config.boot_partition_size
    ], check=True)
    # TODO: The full size of the partitioning will not always be 100%
    full_size = '100%'  # or, e.g. '100G'
    subprocess.run([
        'sudo', 'parted', '-s', target_disk,
        'mkpart', 'primary', config.boot_partition_size, full_size
    ], check=True)

    # TODO: On a symmetrical setup (which should be standard on hardware),
    # we need to repeat the above for the secondary disk.
    # TODO: Needs an option for symmetrical setup (second disk).

    # TODO: Need an option for EFI boot or not.
    # For EFI: sudo parted -s /dev/nbd0 set 1 esp on
    # Up to now, we get by with legacy boot, but some hardware may not
    # support that.
    subprocess.run([
        'sudo', 'parted', '-s', target_disk,
        'set', '1', 'boot', 'on'
    ], check=True)
    subprocess.run([
        'sudo', 'parted', '-s', target_disk,
        'set', '2', 'lvm', 'on'
    ], check=True)

    print("Boot file system")
    subprocess.run(['sudo', 'partprobe'], check=True)
    # TODO: The partitions will be named differently on different devices.
    subprocess.run([
        'sudo', 'mkfs.ext4', '-L', 'boot', '-m', '0', '/dev/nbd0p1'
    ], check=True, stdout=subprocess.DEVNULL)
    # For EFI: mkfs.vfat -n EFI /dev/nbd0p1

    # TODO: The LVM should take care of redundancy (RAID level 1) in
    # symmetrical setups.

    # TODO: Hardware setups may need DRBD volumes for synchronized data
    # storage. This should not be handled by this script in the first
    # iteration, because the OS boot must not depend on such volumes.
    # They can be installed on free sections of the hard drives later on.
    # In a later step, allow for an option to build a one-legged DRBD setup.
    # Each DRBD requires its own volume group, so this part of the config must
    # be additional to the current config anyway.
    # Many services should *not* use DRBD. All services which provide
    # synchronization or replication should use built-in methods. Examples:
    # postgresql (has replication), DHCP (has replication), LDAP (has
    # replication), BIND9 (has replication). Good uses of DRBD are shares and
    # cloud storage services. Note that when using DRBD, use drbd9 and
    # drbdmanage to be prepared for larger clusters.

    print("Make LVM group")
    subprocess.run(['sudo', 'pvcreate', '/dev/nbd0p2'], check=True)
    subprocess.run([
        'sudo', 'vgcreate', config.volgroup, '/dev/nbd0p2'
    ], check=True)

    print("Create volumes and file systems")
    for volume in config.volumes:
        subprocess.run([
            'sudo', 'lvcreate',
            '--size', volume['size'], '--name', volume['name'],
            config.volgroup], check=True)

        if volume['name'] == 'swap':
            subprocess.run(
                ['sudo', 'mkswap',
                 '/dev/{}/{}'.format(config.volgroup, volume['name'])
                 ],
                check=True
            )
            subprocess.run(
                ['sudo', 'swaplabel', '-L', 'swap',
                 '/dev/{}/{}'.format(config.volgroup, volume['name'])
                 ],
                check=True
            )
        else:
            subprocess.run([
                'sudo', 'mkfs.ext4',
                '/dev/{}/{}'.format(config.volgroup, volume['name']),
                '-L', volume['name']], check=True, stdout=subprocess.DEVNULL)


def make_blockdev():
    "Mount image into block device"
    print("Mount image into block device")
    subprocess.run(['sudo', 'modprobe', 'nbd'], check=True)
    subprocess.run([
        'sudo', 'qemu-nbd', '--connect=/dev/nbd0',
        '--persistent', '--fork', '--shared=4', image_filename
    ], check=True)


def destroy_blockdev():
    "Remove loopback device"
    print("Remove the loopback device")
    cmds = [
        ['qemu-nbd', '--disconnect', '/dev/nbd0'],
        ['partprobe'],
        ['rmmod', 'nbd'],
    ]

    for cmd in cmds:
        subprocess.run(['sudo'] + cmd, check=False)


def mount_system():
    "Mount file systems"
    print("Mount file systems")
    subprocess.run(['sudo', 'partprobe'])
    subprocess.run(['sudo', 'vgchange', '-a', 'y', config.volgroup])
    for mountpoint, device in mountpoints:
        directory = os.path.join(mount_path, mountpoint.lstrip('/'))
        subprocess.run([
            'sudo', 'mkdir', '-p', directory
        ], check=True)
        subprocess.run([
            'sudo', 'mount', device, directory
        ], check=True)


def unmount_system(zerofree=False):
    "Reverse all setup procedures and return system back to normal"
    # It can also be used to repair an interrupted generation procedure.
    # This also optionally zeroes the partitions it unmounts.

    # If the option "zerofree" is set, zerofree is called before the volume
    # group is destroyed.'''

    print("Unmount system")
    for mountpoint, device in reversed(mountpoints):
        subprocess.run([
            'sudo', 'umount', device
        ], stderr=subprocess.DEVNULL, check=False)

    if zerofree:
        zerofree_system()

    print("Remove LVM volumes")
    subprocess.run([
        'sudo', 'vgchange', '-a', 'n', config.volgroup
    ], check=False)


def zerofree_system():
    "Run zerofree on all partitions."
    print("Zerofree file systems")
    for mountpoint, device in mountpoints:
        subprocess.run([
            'sudo', 'zerofree', '-v', device
        ], check=False)


def unpack_system():
    "Unpack archive into freshly created and mounted image"
    # CHECK: is everything set up?
    cmd = ['tar', '--extract', '--preserve-permissions',
           '--same-permissions', '--numeric-owner',
           ]
    print("Unpack archive")
    _sudo(cmd + [
        '--directory', mount_path,
        '--file', os.path.join(base_path, config.template_archive),
    ])
    print("Unpack database")
    target = f'{mount_path}/vol/postgresql/var/lib/postgresql/migration'
    _sudo(['mkdir', '-p', target])
    _sudo(cmd + [
        '--directory', target,
        '--file', f'{base_path}/{config.template_archive_db}',
    ])
    # Create symlink for pg_work
    if 'pg_work' in {vol['name'] for vol in config.volumes}:
        _sudo(['ln', '-s', '/vol/pg_work',
               f'{mount_path}/vol/postgresql/var/lib/postgresql/pg_work'])


def install_grub():
    "Install grub"
    print("Install grub in the chroot environment")
    _chroot_cmds(
        ['update-grub'],
        ['grub-install', '/dev/nbd0']
    )


def mount_chroot():
    '''Add mountpoints for chrooting into the system.'''
    print("Mount devices for chroot")
    for mountpoint in chroot_mountpoints:
        directory = os.path.join(mount_path, mountpoint.lstrip('/'))
        subprocess.run([
            'sudo', 'mount', '-o', 'bind', mountpoint, directory
        ], check=True)
    # Transfer resolv.conf over to the chroot
    subprocess.run([
        'sudo', 'mv',
        os.path.join(mount_path, 'etc/resolv.conf'),
        os.path.join(mount_path, 'etc/resolv.conf-chroot'),
    ], check=True)
    subprocess.run([
        'sudo', 'cp',
        '/etc/resolv.conf',
        os.path.join(mount_path, 'etc/resolv.conf'),
    ], check=True)


def unmount_chroot():
    '''Remove system mount points'''
    print("Unmount devices for chroot")
    # Restore resolv.conf to the original
    subprocess.run([
        'sudo', 'mv',
        os.path.join(mount_path, 'etc/resolv.conf-chroot'),
        os.path.join(mount_path, 'etc/resolv.conf'),
    ], check=False)

    for mountpoint in reversed(chroot_mountpoints):
        directory = os.path.join(mount_path, mountpoint.lstrip('/'))
        subprocess.run([
            'sudo', 'umount', directory
        ], stderr=subprocess.DEVNULL, check=False)


def fixup_postgresql(cluster='main'):
    "Move postgresql restore directory into proper place"
    "Returns the data folder of the new database."

    # Move backup from source into correct data dir
    source = '/vol/postgresql/var/lib/postgresql/migration'
    pg_version = subprocess.run(
        ['sudo', 'cat', f'{mount_path}/{source}/PG_VERSION'],
        stdout=subprocess.PIPE,
        universal_newlines=True,
        check=True,
    ).stdout.strip()
    pg_home = '/vol/postgresql/var/lib/postgresql'
    target = f'{pg_home}/{pg_version}'
    pg_data = f'{target}/{cluster}'
    try:
        os.stat(f'{mount_path}/{pg_data}')
    except OSError:
        pass
    else:
        print("PostgreSQL directory found. Bailing out")
        return
    print("Perform postgresql fixup")
    _chroot_cmds(
        ['mkdir', '-p', target],
        ['mv', source, pg_data],
        ['chown', 'postgres:postgres', pg_home, target, pg_data],
        ['chmod', '0700', pg_data],
    )
    if cluster == 'main':
        # Only the main cluster uses a symlink for pg_wal
        pgwal = f'{pg_data}/pg_wal'
        pgwaldir = f'/vol/pg_wal/{pg_version}'
        pgwaltgt = f'{pgwaldir}/main'

        _chroot_cmds(
            ['mkdir', '-p', pgwaldir],
            ['test', '!', '-d', pgwaltgt],  # error if already exists
            ['mv', pgwal, pgwaltgt],
            ['ln', '-s', pgwaltgt, pgwal],
        )

    return pg_data


def create_backup_folders():
    '''
    Create subfolders in /vol/backup. Requires chroot.
    '''
    print("Create folders in /vol/backup")
    for user, folder in [
            ('postgres', 'postgresql'),
            ('root', 'system'),
            ('zope', 'zope'),
    ]:
        target = '/vol/backup/{}'.format(folder)
        _chroot_cmds(
            ['mkdir', '-m', '0770', target],
            ['chown', '{}:backup'.format(user), target],
        )


def unpack_source_data():
    '''
    Unpack source archive into chroot.
    '''
    archive = getattr(config, 'source_archive', None)
    if archive is None:
        return
    if not os.path.exists(archive):
        print('Warning: source_archive set but not present. Skipping'
              ' extraction. Have you forgotten to run fetch-archive for the'
              ' source system?')
        return

    print('Unpacking source data')
    target = f'{mount_path}/opt/perfact/custom/migration/source_system'
    subprocess.run(['sudo', 'mkdir', '-p', target], check=True)
    subprocess.run(
        ['sudo', 'tar', 'xf', archive, '-C', target],
        check=True,
    )


def _prepare_datafs_access():
    '''
    1. Change permissions to Data.FS so user perfact can use
    zoperecord/zopeplayback.
    2. Create copy of Zope config file that accesses the Data.FS directly
    (without ZEO).
    3. Create copy of zodbsync config file that uses the file from 2.
    '''
    zeovar = '/var/lib/zope2.13/zeo/emazeo/var'
    _chroot_cmds(
        ['chgrp', '-R', 'perfact', zeovar, ],
        ['chmod', '-R', 'g+w', zeovar, ],
    )

    # Lines between <zeoclient> and </zeoclient> should be dropped.
    # Instead, we will insert lines with <filestorage>.
    drop = False
    lines = []
    zopeconf = f'{mount_path}/var/lib/zope2.13/instance/ema/etc/zope.conf'
    for line in open(zopeconf).readlines():
        if line.strip() == '<zeoclient>':
            drop = True
            continue

        if line.strip() == '</zeoclient>':
            drop = False
            lines.append('''
              <filestorage>
                path /var/lib/zope2.13/zeo/emazeo/var/Data.fs
              </filestorage>
            ''')
            continue

        if not drop:
            lines.append(line)

    tmpdir = '/opt/perfact/custom/migration/tmp'
    subprocess.run(['sudo', 'mkdir', '-p', f'{mount_path}/{tmpdir}'],
                   check=True)
    subprocess.run(
        ['sudo', 'tee', f'{mount_path}/{tmpdir}/zope.conf'],
        stdout=subprocess.DEVNULL,
        text=True,
        input=''.join(lines)
    )

    lines = []
    for line in open(f'{mount_path}/etc/perfact/modsync/zodb.py').readlines():
        if line.startswith('conf_path = '):
            line = f"conf_path = '{tmpdir}/zope.conf'\n"
        lines.append(line)

    subprocess.run(
        ['sudo', 'tee', f'{mount_path}/{tmpdir}/zodb.conf'],
        stdout=subprocess.DEVNULL,
        text=True,
        input=''.join(lines)
    )


def upload_source_datafs():
    '''Include Data.FS of source system in target under /PerFact/OLD.'''
    _prepare_datafs_access()

    migration = '/opt/perfact/custom/migration'
    repo = '/opt/perfact/dbutils-zoperepo'
    extensions = '/var/lib/zope2.13/instance/ema/Extensions'

    _chroot_cmds(
        ['mv',
         f'{migration}/source_system/{repo}/__root__',
         f'{repo}/__root__/PerFact/OLD'
         ],
        ['mv', f'{migration}/source_system/{extensions}',
         f'{extensions}'],  # TODO: glob files
        ['sudo', '-u', 'perfact',
         'perfact-zopeplayback', '-c', f'{migration}/tmp/zodb.conf',
         '/PerFact/OLD',
         ],
    )


def unpack_source_db():
    '''
    Unpack the source db and set it up to serve port 5433.
    '''
    archive = getattr(config, 'source_archive_db', None)
    if archive is None:
        return
    if not os.path.exists(archive):
        print('Warning: source_archive_db set but not present. Skipping'
              ' extraction. Have you forgotten to run fetch-archive for the'
              ' source system?')
        return

    print('Unpacking source database')
    _sudo(
        ['tar', 'xf', config.source_archive_db,
         '-C', f'{mount_path}/vol/postgresql/var/lib/postgresql/migration'],
    )
    datadir = fixup_postgresql(cluster='source')
    _chroot_cmds(
        ['/opt/perfact/migration/src/inchroot/setup-source-cluster', datadir]
    )


def update_migration_tools():
    '''
    Push the current state of this repository into the chroot.
    '''
    current_user = getpass.getuser()
    path_in_chroot = '/opt/perfact/migration'
    remote_name = 'build-targetmnt'
    tmp_branch = 'temp'
    path = f'{mount_path}{path_in_chroot}'
    assert os.path.exists(path), \
        'Tool repository not found in template.'

    # current and in-chroot git commands
    gitsrc = ['git', '-C', env.base_path]
    gittgt = ['git', '-C', path]
    current_branch = subprocess.check_output(
        gitsrc + ['rev-parse', '--abbrev-ref', 'HEAD'],
        universal_newlines=True,
    ).strip()

    # change owner of inner repo to current user
    _run(['sudo', 'chown', '-R', current_user+':'+current_user, path])
    # add remote to outer repo
    _run(gitsrc + ['remote', 'rm', remote_name], check=False)
    _run(gitsrc + ['remote', 'add', remote_name, path])
    # check out temp branch in inner repo
    _run(gittgt + ['reset', '--hard'])
    _run(gittgt + ['clean', '-df'])
    _run(gittgt + ['checkout', '-B', tmp_branch])
    # force push to inner
    _run(gitsrc + ['push', '--force', remote_name, f'HEAD:{current_branch}'])
    # check out the pushed branch in inner repo
    _run(gittgt + ['checkout', current_branch])
    # delete temp branch in inner repo
    _run(gittgt + ['branch', '-D', tmp_branch])
    # change owner back to perfact
    _chroot_cmds(
        ['chown', '-R', 'perfact:perfact', path_in_chroot]
    )


def clone_custom_files():
    '''
    Clone the given custom file repository into the target system.
    '''
    repo = getattr(config, 'custom_files_repo', None)
    if repo is None:
        return
    subprocess.run(
        ['sudo', 'git', 'clone', repo,
         f'{mount_path}/opt/perfact/custom/migration/files',
         ],
        check=True,
    )


def convert_disk():
    "Convert disk into VM format"
    if config.image_file_vm_format is None:
        return
    print('Exporting image to target format')
    subprocess.run([
        'qemu-img', 'convert', '-p', '-O', config.image_file_vm_format,
        image_filename, os.path.join(build_path, config.image_file_vm),
    ], check=True)