#!/usr/bin/env python3

# Copyright (c) 2021, AT&T Intellectual Property.
# All rights reserved.
#
# SPDX-License-Identifier: LGPL-2.1-only

"""This updates the CPU affinity settings for the dataplane."""

import sys
import configparser
from argparse import ArgumentParser

DP_CONF = '/etc/vyatta/dataplane.conf'


def _delete_cpumask(dpcfg, key_cpumask='cpumask'):

    if not dpcfg.has_option('Dataplane', key_cpumask):
        return

    with open(DP_CONF, 'w') as dpcfg_fp:
        del dpcfg['Dataplane'][key_cpumask]
        dpcfg.write(dpcfg_fp, space_around_delimiters=False)


def _set_cpumask(dpcfg, cpumask, key_cpumask='cpumask'):

    if dpcfg['Dataplane'].get(key_cpumask) == cpumask:
        return

    with open(DP_CONF, 'w') as dpcfg_fp:
        dpcfg['Dataplane'][key_cpumask] = cpumask
        dpcfg.write(dpcfg_fp, space_around_delimiters=False)


def _parse_range(range_str):
    cpuset = set()
    for token in range_str.split(','):
        if '-' in token:
            start, stop = token.split('-')
            for cpuid in range(int(start), int(stop) + 1):
                cpuset.add(cpuid)
        elif token:
            cpuset.add(int(token))

    return cpuset


def _validate_cpumask(cpumask, control_cpumask):
    try:
        cpuset = _parse_range(cpumask)
        control_cpuset = _parse_range(control_cpumask)
    except ValueError as err:
        print("Validation of CPU affinity value failed : ", err)
        sys.exit(1)

    # limited by CPU_SETSIZE (sched.h) and pthread_setaffinity_np
    if max(control_cpuset) >= 1024:
        print("Dataplane Control CPU affinity exceeds number of "
              "supported CPUs.")
        sys.exit(1)

    if not control_cpuset.issubset(cpuset):
        print("Dataplane Control CPU affinity must be a subset of "
              "the general Dataplane CPU affinity.")
        sys.exit(1)


def _main():
    key_cpumask = 'cpumask'

    arg_parser = ArgumentParser()
    arg_parser.add_argument('--set', action='store', metavar='CPUMASK',
                            help='CPU range', required=False)
    arg_parser.add_argument('--delete', action='store_true',
                            help='Delete CPU affinity config', required=False)
    arg_parser.add_argument('--control', action='store_true',
                            help='Perfrom update of CPU affinity config for control threads',
                            required=False)
    arg_parser.add_argument('--validate', nargs='+', action='store',
                            metavar=('CPUMASK', 'CONTROL_CPUMASK'),
                            help='Validate CPU affinity settings', required=False)

    args = arg_parser.parse_args()

    cpumask = args.set
    if args.control:
        key_cpumask = 'control_cpumask'

    if not args.delete and not args.set and not args.validate:
        print('No action specified.')
        sys.exit(1)

    dpcfg = configparser.ConfigParser()
    try:
        dpcfg.read(DP_CONF)
    except configparser.Error as err:
        print('Failed to load {}: {}'.format(DP_CONF, err))
        sys.exit(1)

    if dpcfg.has_section('Dataplane') is False:
        print('Invalid dataplane config. Aborting.')
        sys.exit(1)

    if args.delete:
        try:
            _delete_cpumask(dpcfg, key_cpumask)
        except (IOError, configparser.Error) as err:
            print('Failed to delete CPU affinity setting from dataplane configuration: {}'
                  .format(err))
            sys.exit(1)
    elif args.validate:
        if len(args.validate) != 2:
            print('Failed to verify dataplane CPU affinity.')
            sys.exit(1)
        try:
            _validate_cpumask(args.validate[0], args.validate[1])
        except (IOError, configparser.Error) as err:
            print('Failed to validate CPU affinity settings: {}'
                  .format(err))
            sys.exit(1)
    else:
        try:
            _set_cpumask(dpcfg, cpumask, key_cpumask)
        except (IOError, configparser.Error) as err:
            print('Failed to update CPU affinity in dataplane configuration: {}'.format(err))
            sys.exit(1)


if __name__ == '__main__':
    _main()
