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

#############################################################################
#
# This is a wrapper script that can be used to power cycle a system
#

try:
    import os
    import subprocess
    import time
    import sys
    import logging
    import logging.handlers
except ImportError as e:
    raise ImportError(str(e) + "- required module not found")    

# initialize the global logger
SYSLOG_IDENTIFIER = "cl-powercycle"
global logger
logger = logging.getLogger(SYSLOG_IDENTIFIER)
fmt = logging.Formatter(fmt='%(name)s %(levelname)s: %(message)s')
handler = logging.handlers.SysLogHandler('/dev/log')
handler.setFormatter(fmt)
logger.setLevel(logging.INFO)
#logger.setLevel(logging.DEBUG)
logger.addHandler(handler)

# Define paths
reset_thermal_asic_path = "/var/run/hw-management/system/reset_asic_thermal"
shut_unlock_path = "/var/run/hw-management/system/shutdown_unlock"
pwrcycle_path = "/var/run/hw-management/system/pwr_cycle"

# Check for -noprompt option
noprompt = False
if len(sys.argv) > 1 and sys.argv[1] == '-noprompt':
    noprompt = True

# Check if system was rebooted due to ASIC thermal shutdown
if os.path.exists(reset_thermal_asic_path):
    try:
        with open(reset_thermal_asic_path, 'r') as f:
            reset_value = f.read().strip()
            if reset_value == '1':
                # Log that trying to recover thermal asic shut case
                logger.debug("Trying to recover from thermal ASIC shutdown...")

                # Release system from thermal shutdown lock and power cycle
                if os.path.exists(shut_unlock_path):
                    try:
                        with open(shut_unlock_path, 'w') as unlock_file:
                            unlock_file.write('1')
                            logger.info("thermal asic shutdown lock released.")
                    except IOError as e:
                        logger.error("Error writing to thermal unlock file:", e)
                        exit(1)
                else:
                    logger.error("Thermal asic unlock is not supported")
            else:
                logger.debug("System was not rebooted due to ASIC thermal shutdown.")
    except Exception as e:
        logger.error("An error occurred checking thermal asic path:", e)
else:
    logger.debug("Thermal asic file not found:", reset_thermal_asic_path)

logger.info("This script will now power cycle the switch.")
print("This script will now power cycle the switch.")

# Ask user if they want to proceed with shutdown (unless -noprompt is used)
if not noprompt:
    response = input("Do you want to proceed with power-cycle? (yes/no): ")
    if response.lower() != "yes":
        logger.info("Switch power-cycle aborted.")
        print("Switch power-cycle aborted.")
        exit(0)

try:
    with open(pwrcycle_path, 'w') as f:
        logger.critical("Power cycling the switch ...")
        subprocess.run(["wall", "System is undergoing a power cycle ... "])
        handler.flush()
        handler.close()
        # Force system to write all pending data to disk
        subprocess.run(["sync"], check=True)
        time.sleep(1)
        f.write('1')
except Exception as e:
    logger.error("Error writing to pwr_cycle file :", e)
    exit(1)

