#!/usr/bin/python3
#
# Copyright 2018 ScyllaDB
#

#
# This file is part of Scylla.
#
# Scylla is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# Scylla is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with Scylla.  If not, see <http://www.gnu.org/licenses/>.

import os
import argparse
import pwd
import grp
import sys
import stat
from scylla_util import *

if __name__ == '__main__':
    if os.getuid() > 0:
        print('Requires root permission.')
        sys.exit(1)
    parser = argparse.ArgumentParser(description='Configure RAID volume for Scylla.')
    parser.add_argument('--disks', required=True,
                        help='specify disks for RAID')
    parser.add_argument('--raiddev', default='/dev/md0',
                        help='MD device name for RAID')
    parser.add_argument('--update-fstab', action='store_true', default=False,
                        help='update /etc/fstab for RAID')
    parser.add_argument('--root', default='/var/lib/scylla',
                        help='specify the root of the tree')
    parser.add_argument('--volume-role', default='all',
                        help='specify how will this device be used (data, commitlog, or all)')
    parser.add_argument('--force-raid', action='store_true', default=False,
                        help='force constructing RAID when only one disk is specified')

    args = parser.parse_args()

    root = args.root.rstrip('/')
    if args.volume_role == 'all':
        mount_at=root
    elif args.volume_role == 'data':
        mount_at='{}/data'.format(root)
    elif args.volume_role == 'commitlog':
        mount_at='{}/commitlog'.format(root)
    else:
        print('Invalid role specified ({})'.format(args.volume_role))
        parser.print_help()
        sys.exit(1)

    disks = args.disks.split(',')
    for disk in disks:
        if not os.path.exists(disk):
            print('{} is not found'.format(disk))
            sys.exit(1)
        if not stat.S_ISBLK(os.stat(disk).st_mode):
            print('{} is not block device'.format(disk))
            sys.exit(1)
        if not is_unused_disk(disk):
            print('{} is busy'.format(disk))
            sys.exit(1)

    if os.path.exists(args.raiddev):
        print('{} is already using'.format(args.raiddev))
        sys.exit(1)

    if os.path.ismount(mount_at):
        print('{} is already mounted'.format(mount_at))
        sys.exit(1)

    if is_debian_variant():
        run('env DEBIAN_FRONTEND=noninteractive apt-get -y install mdadm xfsprogs')
    elif is_gentoo_variant():
        run('emerge -uq sys-fs/mdadm sys-fs/xfsprogs')

    if len(disks) == 1 and not args.force_raid:
        raid = False
        fsdev = disks[0]
    else:
        raid = True
        fsdev = args.raiddev

    print('Creating {type} for scylla using {nr_disk} disk(s): {disks}'.format(type='RAID0' if raid else 'XFS volume', nr_disk=len(disks), disks=args.disks))
    if dist_name() == 'Ubuntu' and dist_ver() == '14.04':
        if raid:
            run('udevadm settle')
            run('mdadm --create --verbose --force --run {raid} --level=0 -c1024 --raid-devices={nr_disk} {disks}'.format(raid=fsdev, nr_disk=len(disks), disks=args.disks.replace(',', ' ')))
            run('udevadm settle')
        run('mkfs.xfs {} -f'.format(fsdev))
    else:
        procs=[]
        for disk in disks:
            d = disk.replace('/dev/', '')
            discard_path = '/sys/block/{}/queue/discard_granularity'.format(d)
            if os.path.exists(discard_path):
                with open(discard_path) as f:
                    discard = f.read().strip()
                if discard != '0':
                    proc = subprocess.Popen(['blkdiscard', disk])
                    procs.append(proc)
        for proc in procs:
            proc.wait()
        if raid:
            run('udevadm settle')
            run('mdadm --create --verbose --force --run {raid} --level=0 -c1024 --raid-devices={nr_disk} {disks}'.format(raid=fsdev, nr_disk=len(disks), disks=args.disks.replace(',', ' ')))
            run('udevadm settle')
        run('mkfs.xfs {} -f -K'.format(fsdev))

    if is_debian_variant():
        confpath = '/etc/mdadm/mdadm.conf'
    else:
        confpath = '/etc/mdadm.conf'

    if raid:
        res = out('mdadm --detail --scan')
        with open(confpath, 'w') as f:
            f.write(res)

    makedirs(mount_at)
    run('mount -t xfs -o noatime {raid} "{mount_at}"'.format(raid=fsdev, mount_at=mount_at))

    makedirs('{}/data'.format(root))
    makedirs('{}/commitlog'.format(root))
    makedirs('{}/coredump'.format(root))

    uid = pwd.getpwnam('scylla').pw_uid
    gid = grp.getgrnam('scylla').gr_gid
    os.chown(root, uid, gid)
    os.chown('{}/data'.format(root), uid, gid)
    os.chown('{}/commitlog'.format(root), uid, gid)
    os.chown('{}/coredump'.format(root), uid, gid)

    if args.update_fstab:
        res = out('blkid {}'.format(fsdev))
        match = re.search(r'^/dev/\S+: (UUID="\S+")', res.strip())
        uuid = match.group(1)
        with open('/etc/fstab', 'a') as f:
            f.write('{uuid} {mount_at} xfs noatime,nofail 0 0\n'.format(uuid=uuid, mount_at=mount_at))
        mounts_conf = '/etc/systemd/system/scylla-server.service.d/mounts.conf'
        if not os.path.exists(mounts_conf):
            makedirs('/etc/systemd/system/scylla-server.service.d/')
            with open(mounts_conf, 'w') as f:
                f.write('[Unit]\nRequiresMountsFor={mount_at}\n'.format(mount_at=mount_at))
        else:
            with open(mounts_conf, 'a') as f:
                f.write('RequiresMountsFor={mount_at}\n'.format(mount_at=mount_at))

    if is_debian_variant():
        run('update-initramfs -u')
