#!/usr/bin/python
# Copyright 2016, 2017, Cumulus Networks, Inc.  All rights reserved.

import json
import sys
import os
import argparse

class ArgParseError(RuntimeError):
    pass

if __name__ == '__main__':
    if (os.geteuid() != 0) :
        sys.stderr.write('root privileges are needed to run riot-update\n')
        sys.exit(-1)

    parser = argparse.ArgumentParser(
        description='Update cumulus RIOT resource allocations')
    parser.add_argument('-v', '--verbose',
                        required=False,
                        action='store_true',
                        help='Verbose output')
    parser.add_argument('-p', '--profile',
                        required=False,
                        default='default',
                        help='RIOT profile')
    try:
        args = parser.parse_args()
    except ArgParseError, e:
        parser.error(str(e))

    riot_file = '/etc/bcm.d/datapath/riot.json'
    riot_bcm_file = '/etc/bcm.d/config.d/04vxlan_riot.bcm'

    if not os.path.isfile(riot_file):
        if args.verbose:
            sys.stdout.write('Skip riot setup of %s; riot.json absent\n' % args.profile)
        sys.exit(0)

    riot_db = {}
    with open(riot_file, 'r') as fp:
        riot_db = json.load(fp)

    profile = args.profile

    if profile == 'disable':
        if os.path.isfile(riot_bcm_file):
            os.remove (riot_bcm_file)
        sys.exit(0)

    profile_elem = riot_db.get(profile)
    if not profile_elem:
        profile = 'default'
        profile_elem = riot_db.get(profile)

    config = []

    if args.verbose:
        sys.stdout.write('Setting up the RIOT for %s\n' % args.profile)

    config.append('riot_enable=1\n')
    config.append('riot_overlay_l3_intf_mem_alloc_mode=1\n')
    config.append('riot_overlay_l3_egress_mem_alloc_mode=1\n')

    riot_overlay_l3_intf_mem_size = profile_elem.get('riot_overlay_l3_intf_mem_size', 0)
    if riot_overlay_l3_intf_mem_size:
        config.append('riot_overlay_l3_intf_mem_size=%d\n' % riot_overlay_l3_intf_mem_size)

    riot_overlay_l3_egress_mem_size = profile_elem.get('riot_overlay_l3_egress_mem_size', 0)
    if riot_overlay_l3_egress_mem_size:
        config.append('riot_overlay_l3_egress_mem_size=%d\n' % riot_overlay_l3_egress_mem_size)

    config.append('l3_ecmp_levels=2\n')
    config.append('riot_overlay_ecmp_resilient_hash_size=16384\n')
    if config:
        with open (riot_bcm_file, 'w')as fp:
            fp.write(''.join(config))
