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

import json
import sys
import os
import argparse

import cumulus.platforms

class ArgParseError(RuntimeError):
    pass

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

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

    uft_file = '/etc/bcm.d/datapath/uft.json'
    uft_bcm_file = '/etc/bcm.d/config.d/02uft.bcm'

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

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

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

    profile = args.profile
    profile_elem = uft_db.get(profile)
    if profile_elem:
        if profile_elem.get('useDefault'):
            profile = 'default'
            profile_elem = uft_db.get(profile)
    else:
        sys.stderr.write('UFT profile does not exist using default\n')
        profile = 'default'
        profile_elem = uft_db.get(profile)

    if not profile_elem:
        if args.verbose:
            sys.stdout.write('Skipped UFT setup; no %s or %s\n' % (args.profile, profile))
        sys.exit(0)

    config = []

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

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

    alpmEnable = profile_elem.get('alpmEnable', False)
    if alpmEnable:
        config.append('l3_alpm_enable=2\n')
        config.append('ipv6_lpm_128b_enable=1\n')

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

    platform = cumulus.platforms.probe()
    chip = platform.switch.chip
    if (isinstance(chip, cumulus.platform.TridentThreeX7Chip) or \
        isinstance(chip, cumulus.platform.TridentThreeX7_56870_Chip) or \
        isinstance(chip, cumulus.platform.TridentThreeX7_56873_Chip)):
        if alpmEnable:
	    config.append('fpem_mem_entries=%d\n' % 32768)

    if config:
        with open (uft_bcm_file, 'w')as fp:
            fp.write(''.join(config))
