#
# upload_extra_script.py
# set the output_port
#  if target_filename is found then that drive is used
#  else if target_drive is found then that drive is used
#
from __future__ import print_function

import pioutil
if pioutil.is_pio_build():

	target_filename = "FIRMWARE.CUR"
	target_drive = "REARM"

	import os,getpass,platform

	current_OS = platform.system()
	Import("env")

	def print_error(e):
		print('\nUnable to find destination disk (%s)\n' \
			  'Please select it in platformio.ini using the upload_port keyword ' \
			  '(https://docs.platformio.org/en/latest/projectconf/section_env_upload.html) ' \
			  'or copy the firmware (.pio/build/%s/firmware.bin) manually to the appropriate disk\n' \
			  %(e, env.get('PIOENV')))

	def before_upload(source, target, env):
		try:
			#
			# Find a disk for upload
			#
			upload_disk = 'Disk not found'
			target_file_found = False
			target_drive_found = False
			if current_OS == 'Windows':
				#
				# platformio.ini will accept this for a Windows upload port designation: 'upload_port = L:'
				#   Windows - doesn't care about the disk's name, only cares about the drive letter
				import subprocess,string
				from ctypes import windll

				# getting list of drives
				# https://stackoverflow.com/questions/827371/is-there-a-way-to-list-all-the-available-drive-letters-in-python
				drives = []
				bitmask = windll.kernel32.GetLogicalDrives()
				for letter in string.ascii_uppercase:
					if bitmask & 1:
						drives.append(letter)
					bitmask >>= 1

				for drive in drives:
					final_drive_name = drive + ':\\'
					# print ('disc check: {}'.format(final_drive_name))
					try:
						volume_info = str(subprocess.check_output('cmd /C dir ' + final_drive_name, stderr=subprocess.STDOUT))
					except Exception as e:
						print ('error:{}'.format(e))
						continue
					else:
						if target_drive in volume_info and not target_file_found:  # set upload if not found target file yet
							target_drive_found = True
							upload_disk = final_drive_name
						if target_filename in volume_info:
							if not target_file_found:
								upload_disk = final_drive_name
							target_file_found = True

			elif current_OS == 'Linux':
				#
				# platformio.ini will accept this for a Linux upload port designation: 'upload_port = /media/media_name/drive'
				#
				drives = os.listdir(os.path.join(os.sep, 'media', getpass.getuser()))
				if target_drive in drives:  # If target drive is found, use it.
					target_drive_found = True
					upload_disk = os.path.join(os.sep, 'media', getpass.getuser(), target_drive) + os.sep
				else:
					for drive in drives:
						try:
							files = os.listdir(os.path.join(os.sep, 'media', getpass.getuser(), drive))
						except:
							continue
						else:
							if target_filename in files:
								upload_disk = os.path.join(os.sep, 'media', getpass.getuser(), drive) + os.sep
								target_file_found = True
								break
				#
				# set upload_port to drive if found
				#

				if target_file_found or target_drive_found:
					env.Replace(
						UPLOAD_FLAGS="-P$UPLOAD_PORT"
					)

			elif current_OS == 'Darwin':  # MAC
				#
				# platformio.ini will accept this for a OSX upload port designation: 'upload_port = /media/media_name/drive'
				#
				drives = os.listdir('/Volumes')  # human readable names
				if target_drive in drives and not target_file_found:  # set upload if not found target file yet
					target_drive_found = True
					upload_disk = '/Volumes/' + target_drive + '/'
				for drive in drives:
					try:
						filenames = os.listdir('/Volumes/' + drive + '/')   # will get an error if the drive is protected
					except:
						continue
					else:
						if target_filename in filenames:
							if not target_file_found:
								upload_disk = '/Volumes/' + drive + '/'
							target_file_found = True

			#
			# Set upload_port to drive if found
			#
			if target_file_found or target_drive_found:
				env.Replace(UPLOAD_PORT=upload_disk)
				print('\nUpload disk: ', upload_disk, '\n')
			else:
				print_error('Autodetect Error')

		except Exception as e:
			print_error(str(e))

	env.AddPreAction("upload", before_upload)