#!/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 sys
import argparse
import subprocess
import re
from scylla_util import *

def bool2str(val):
    return 'yes' if val else 'no'

def str2bool(val):
    return True if val == 'yes' else False

if __name__ == '__main__':
    if os.getuid() > 0:
        print('Requires root permission.')
        sys.exit(1)
    if is_redhat_variant():
        cfg = sysconfig_parser('/etc/sysconfig/scylla-server')
    else:
        cfg = sysconfig_parser('/etc/default/scylla-server')
    set_nic = str2bool(cfg.get('SET_NIC'))
    ami = str2bool(cfg.get('AMI'))

    parser = argparse.ArgumentParser(description='Setting parameters on Scylla sysconfig file.')
    parser.add_argument('--nic',
                        help='specify NIC')
    parser.add_argument('--mode',
                        help='network mode (posix, dpdk)')
    parser.add_argument('--nr-hugepages', type=int,
                        help='number of hugepages')
    parser.add_argument('--user',
                        help='user (dpdk requires root)')
    parser.add_argument('--group',
                        help='group (dpdk requires root)')
    parser.add_argument('--homedir',
                        help='scylla home directory')
    parser.add_argument('--confdir',
                        help='scylla config directory')
    parser.add_argument('--setup-nic', action='store_true', default=set_nic,
                        help='setup NIC\'s interrupts, RPS, XPS')
    parser.add_argument('--ami', action='store_true', default=ami,
                        help='AMI instance mode')
    args = parser.parse_args()

    if args.nic and not is_valid_nic(args.nic):
            print('NIC {} not found.'.format(args.nic))
            sys.exit(1)

    ifname = args.nic if args.nic else cfg.get('IFNAME')
    network_mode = args.mode if args.mode else cfg.get('NETWORK_MODE')

    if args.setup_nic:
        rps_cpus = out('/usr/lib/scylla/posix_net_conf.sh --cpu-mask {}'.format(ifname))
        if len(rps_cpus) > 0:
            cpuset = hex2list(rps_cpus)
            run('/usr/lib/scylla/scylla_cpuset_setup --cpuset {}'.format(cpuset))

    ethdrv = ''
    ethpciid = ''
    if network_mode == 'dpdk':
        dpdk_status = out('/usr/lib/scylla/dpdk-devbind.py --status')
        match = re.search('if={} drv=(\S+)'.format(ifname), dpdk_status, flags=re.MULTILINE)
        ethdrv = match.group(1)
        match = re.search('^(\\S+:\\S+:\\S+\.\\S+) [^\n]+ if={} '.format(ifname), dpdk_status, flags=re.MULTILINE)
        ethpciid = match.group(1)

    if args.mode:
        cfg.set('NETWORK_MODE', args.mode)
    if args.nic:
        cfg.set('IFNAME', args.nic)
    if cfg.get('ETHDRV') != ethdrv:
        cfg.set('ETHDRV', ethdrv)
    if cfg.get('ETHPCIID') != ethpciid:
        cfg.set('ETHPCIID', ethpciid)
    if args.nr_hugepages:
        cfg.set('NR_HUGEPAGES', args.nr_hugepages)
    if args.user:
        cfg.set('USER', args.user)
    if args.group:
        cfg.set('GROUP', args.group)
    if args.homedir:
        cfg.set('SCYLLA_HOME', args.homedir)
    if args.confdir:
        cfg.set('SCYLLA_CONF', args.confdir)
    if str2bool(cfg.get('SET_NIC')) != args.setup_nic:
        cfg.set('SET_NIC', bool2str(args.setup_nic))
    if str2bool(cfg.get('AMI')) != args.ami:
        cfg.set('AMI', bool2str(args.ami))
    cfg.commit()