import optparse
import subprocess
import cmd
import json

ISDCT_EXE = "./isdct"
EXPECTED_FIRM_REVISIONS = ["XCV10110"]

# Enum class for Status values
class STATUS:
	Unknown = -1
	Success=0
	RMA=1
	FeatureFailed=2
	Cancelled=3
	WrongFirmwareRev = 4
	SMARTPrefail = 5
	WarnSecureErase = 6
	# Status from DCT
	Isdct_FW_Up_To_Date = 64

def extractSerialOfAttachedDrives(output=False):
	# Read the drive inforext
	cmd = subprocess.Popen([ISDCT_EXE, "show", "-intelssd"], stdout=subprocess.PIPE,
							stderr=subprocess.STDOUT)
	cmd.wait()

	driveInfo = cmd.stdout.read().decode()
	if output == True:
		print(driveInfo)

	# Read drive info on all attached drives and then
	driveSerialSplit = driveInfo.split("SerialNumber : ")
	driveSerials = [driveSerialSplit[i].splitlines()[0].strip() for i in range(1, len(driveSerialSplit))]

	return driveSerials

def readLBAs(serial_number, output=False, fullScan=False):
	# Read the LBAs on the drive
	cmd = subprocess.Popen([ISDCT_EXE, "start", "-scan", "ReadScan", "-intelssd", serial_number, "FullScan", "=", str(fullScan)], stdout=subprocess.PIPE,
						   stderr=subprocess.STDOUT)
	cmd.wait()

	if output == True:
		print(cmd.stdout.read().decode())

	if cmd.returncode == STATUS.Success:
		return STATUS.Success

	return STATUS.FeatureFailed

def runSecureErase(serial_number, force=False, output=False):
	# Warn user that secure erase will indeed erase the drive
	proceed = ""
	if force == False:
		proceed = input("Secure erase will wipe ALL data on the drive! Proceed?")

	if force == True or (proceed == 'Y' or proceed == 'y'):
		cmd = subprocess.Popen([ISDCT_EXE, "delete", "-force", "-intelssd", serial_number], stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
		cmd.wait()

		if output == True:
		   print(cmd.stdout.read().decode())

		if cmd.returncode == STATUS.Success:
			return STATUS.Success
	else:
		return STATUS.Cancelled

	return STATUS.FeatureFailed


def runUpdate(serial_number, force=False, output=False):
	# Warn user that secure erase may be performed and that drive should be backed up.
	proceed = ""
	if force == False:
		proceed = input("Firmware update may be performed as part of flow! Proceed?")

	if force == True or (proceed == 'Y' or proceed == 'y'):
		cmd = subprocess.Popen([ISDCT_EXE, "load", "-force", "-intelssd", serial_number], stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
		cmd.wait()

		if output == True:
		   print(cmd.stdout.read().decode())

		if (cmd.returncode == STATUS.Success) or (cmd.returncode == STATUS.Isdct_FW_Up_To_Date):
			return STATUS.Success

	else:
		return STATUS.Cancelled

	return STATUS.FeatureFailed

def checkExpectedRevision(serial_number, output=False):
	# Check the desired firmware revision before returning success
	cmd = subprocess.Popen([ISDCT_EXE, "show", "-intelssd", serial_number], stdout=subprocess.PIPE,
						   stderr=subprocess.STDOUT)
	cmd.wait()

	driveInfo = cmd.stdout.read().decode()
	if output == True:
		print(driveInfo)

	# Check the actual firmware rev.
	driveFirmware = driveInfo.split("Firmware : ")[1].splitlines()[0]
	driveFirmware.strip()

	if driveFirmware in EXPECTED_FIRM_REVISIONS:
		return STATUS.Success

	return STATUS.WrongFirmwareRev

def retrieveSMART(serial_number, output=False, smartId="00", raw=True):
	# Return type to be a named tuple associating status and SMART val (either raw or normalized)
	statusAndVal = [STATUS.Success, -1]

	# First call CLI to get the SMART Attr data
	cmd = subprocess.Popen([ISDCT_EXE, "show","-o", "json", "-smart", smartId, "-intelssd", serial_number], stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
	cmd.wait()

	# Update status portion of the tuple that will be returned
	statusAndVal[0] = cmd.returncode

	if statusAndVal[0] != STATUS.Success:
		return statusAndVal

	smartStr = cmd.stdout.read().decode()
	if output == True:
		print(smartStr)

	try:
		data = json.loads(smartStr)

		if raw == True:
			statusAndVal[1] = int(data["SMART Attributes " + serial_number][smartId]["Raw"])
		else:
			statusAndVal[1] = int(data["SMART Attributes " + serial_number][smartId]["Normalized"])
	except:
		print("Failed to parse the returned SMART JSON!")
		statusAndVal[0] = STATUS.Unknown

	return statusAndVal

def runFlow(serial_number, force=False, output=False):
	result = STATUS.Success

	# Perform a firmware update
	result = runUpdate(serial_number, force, output)

	if result != STATUS.Success:
		return result

	result = checkExpectedRevision(serial_number, output=False)

	if result != STATUS.Success:
		return result

	# Read SMART attr 0x9. Care about raw value.
	result, smart09 = retrieveSMART(serial_number, output, "09" , True)

	if result != STATUS.Success:
		return result

	# <1700 power on hours? Return success.
	if smart09 < 1700:
		return result

	# Read SMART atrr 0xB8. Care about normalized value. Read SMART attr 0xBB. Care about raw value.
	result, smartB8 = retrieveSMART(serial_number, output, "B8" , False)

	if result != STATUS.Success:
		return result

	result, smartBB = retrieveSMART(serial_number, output, "BB", True)

	if result != STATUS.Success:
		return result

	# Is normalized 0xB8 == 100? Then analyze SMART 0xBB. If raw value == 0 then return success.
	if smartB8 == 100 and smartBB == 0:
		return result

	# Else if Normalized 0xB8 value <90? RMA Drive!!! Fatal error.
	if smartB8 != 100 and smartB8 < 90:
		result = STATUS.SMARTPrefail
		return result

	# If we get here then we need to save smart 0xB8 and 0xBB
	oldSmartB8 = smartB8
	oldSmartBB = smartBB

	# Read all LBAs
	result = readLBAs(serial_number, output, True)

	if result != STATUS.Success:
		return result

	# Read SMART atrr 0xB8 again. Care about normalized value. Read SMART attr 0xBB again. Care about raw value.
	result, smartB8 = retrieveSMART(serial_number, output, "B8" , False)

	if result != STATUS.Success:
		return result

	result, smartBB = retrieveSMART(serial_number, output, "BB", True)

	if result != STATUS.Success:
		return result

	# If new SMART attrs <= old SMART attrs, then success.
	if (smartB8 <= oldSmartB8) and (smartBB <= oldSmartBB):
		return result

	# No more secure erase in flow. Warn user to secure erase instead
	# result = runSecureErase(serial_number, force, output)
	result = STATUS.WarnSecureErase
	return result

def optional_arg(arg_default):
	def func(option, opt_str, value, parser):
		if parser.rargs and not parser.rargs[0].startswith('-'):
			val = parser.rargs[0]
			parser.rargs.pop(0)
		else:
			val = 'empty'
		setattr(parser.values, option.dest, val)

	return func

if __name__ == "__main__":
	# Parse the input options
	success = STATUS.Unknown
	parser = optparse.OptionParser()

	parser.add_option('-s',
					  dest='serialNumber',
					  default=None,
					  help='Serial number of drive.')

	parser.add_option('-f',
					  dest='force',
					  default=None,
					  action='callback',
					  callback=optional_arg('empty'),
					  help='Force the flow without prompts.')

	parser.add_option('-o',
					  dest='output',
					  default=None,
					  action='callback',
					  callback=optional_arg('empty'),
					  help='Print all DCT output.')

	parser.add_option('-a',
					  dest='all',
					  default=None,
					  action='callback',
					  callback=optional_arg('empty'),
					  help='Perform special flow on all attached drives.')

	# get command line information
	(options, args) = parser.parse_args()

	force = False
	output = False
	serialNumbers = []

	if options.force != None:
		force = True

	if options.output != None:
		output = True

	# Should flow be run on all attached drives?
	if options.all != None:
		serialNumbers = extractSerialOfAttachedDrives(output)
	# Ensure user provides serial number
	elif options.serialNumber != None:
		serialNumbers.append(options.serialNumber)
	else:
		print("Error: Must supply a serial number value with \'-s\' or \'-a\' option (for all drives).")
		raise SystemExit(-1)

	# Loop all generated serial numbers and run flow
	for serial in serialNumbers:
		success = runFlow(serial, force, output)

		print ("\n================ " + serial + " ================")

		if success == STATUS.Success:
			print ("Success: NAND Channel hang issue has been mitigated.")
		elif success == STATUS.SMARTPrefail:
			print ("Caution (SMART B8h > Threshold): SMART pre-fail condition has been triggered.")
		elif success == STATUS.FeatureFailed:
			print ("Failure: Was not able to run a required feature on the drive. Was not able to finish flow due to driver/feature error. Re-run with -o option for more information.")
		elif success == STATUS.Cancelled:
			print ("Warning: User cancelled flow.")
		elif success == STATUS.WrongFirmwareRev:
			print ("Warning: Incorrect firmware revision for running flow. Either a drive that is not targeted for special flow or a drive not updated to latest firmware. Expected firmware revision(s): " + ", ".join(EXPECTED_FIRM_REVISIONS) + ".")
		elif success == STATUS.WarnSecureErase:
			print ("Notify: Latent read uncorrectable errors exist on this drive and drive will need to be secure erased.")
		else:
			print ("Failure: Failed to fully implement flow. Unrecognized issue occurred.")

		print ("=================" + str(len(serial)* "=" ) + "=================\n")

	SystemExit(success)