#! /usr/bin/python3
# Copyright (C) 2020-2023 NVIDIA Corporation. ALL RIGHTS RESERVED.

#############################################################################
#
# This is a wrapper script that can be used to power down a system
#
try:
    import os
    import subprocess
    import time
except ImportError as e:
    raise ImportError(str(e) + "- required module not found")    

cpu_type_path          = '/run/hw-management/config/cpu_type'
atom_rangeley_cpu_type = 0x64d
broadwell_cpu_type     = 0x656

# Identify the CPU type used on this switch
cpu_type = None
if os.path.exists(cpu_type_path) == True:
    with open(cpu_type_path) as fp:
        val = fp.read().rsplit('\n')[0]
        cpu_type = (int)(val, 16)

# Identify if this switch is an Alligator switch or not
# One could argue that this code is a bit reedundant.
# It has been written this way for the sake of documentation/completeness.
# (ie so there's no doubt that the sn2201_m platform was considered when
# this script was updated.
is_alligator = False
output = subprocess.check_output('/usr/bin/platform-detect').decode()
if "nvidia,sn2201" in output:
    is_alligator = True
elif "nvidia,sn2201_m" in output:
    is_alligator = True

# Linux poweroff does not work on Atom-Rangley, Coffeelake and Alligator systems.
# So we need to implement poweroff in a different way to ensure that the system
# is protected should a thermal event occur.
if cpu_type == atom_rangeley_cpu_type or cpu_type == broadwell_cpu_type or is_alligator:
    # 1st try to make over HALT bit on system where it's supported  
    os.system("echo 1 > /var/run/hw-management/system/pwr_down")
    time.sleep(3)

    # 2nd try to make halt over 2 PSUs
    os.system("echo 1 > /var/run/hw-management/system/psu1_on")
    os.system("echo 1 > /var/run/hw-management/system/psu2_on")
    time.sleep(1)
else:
    subprocess.call("/sbin/poweroff")

