#!/usr/bin/env python3

# Copyright (c) 2019 Arista Networks, Inc.  All rights reserved.
# Arista Networks, Inc. Confidential and Proprietary.

from __future__ import absolute_import, division, print_function
from Ark import getPlatform
import Tac
import re

def getSid():
   with open( '/proc/cmdline', 'r' ) as f:
      m = re.search( r"sid=(\w+)", f.read() )
      if m:
         return m.group( 1 )

   return None

EXCLUDED_PLATFORMS = [
   'baremetal',
   'ceoslab',
   'ceosdsf',
   'dropbear',
   'veos',
   'councilbluffs',
]

EXCLUDED_SID = []

MODPROBE = "/sbin/modprobe"

def modprobeDmamem():
   platform = getPlatform()

   if platform in EXCLUDED_PLATFORMS:
      return

   sid = getSid()
   if sid in EXCLUDED_SID:
      return

   MB = 1024 * 1024
   GB = 1024 * MB

   lo_range_start = 512 * MB
   lo_range_size = 2560 * MB

   hi_range_start = 256 * GB
   hi_range_size = 256 * GB

   if sid in [ "HorseshoeLake", "HorseshoeLakeS" ]:
      lo_range_start = 0

   if platform in [ "redstart", "skylark", "dardanelle" ]:
      hi_range_start = 1024 * GB

   if platform in [ "eaglepeak", "oak" ]:
      hi_range_start = 128 * GB

   # On Sprucefish we are a bit tight on 32bit memory as almost the whole
   # top 2GB of memory are reserved for BAR regions
   if platform in [ "sprucefish" ]:
      lo_range_size = 1024 * MB

   Tac.run( [ MODPROBE, "dmamem_mod", 
             f'iova_lo_range={lo_range_start},{lo_range_size}', 
             f'iova_hi_range={hi_range_start},{hi_range_size}',
             ] )

if __name__ == "__main__":
   modprobeDmamem()
