import inquirer
import time
import os
from artery_loader import FlashAT32
import colorama
from colorama import Fore


class Tester:

    def __init__(self, debug_port):
        self.debug_port = debug_port
        colorama.init(autoreset=True)
    
    def menu(self):
        questions = [
	    inquirer.List('action',
                message="Что делаем",
                choices=['Прошить IAP и FW [MDIO_88]', 
                        'Прошить IAP и FW [MAI_12]', 
                        'Прошить IAP и FW [MAO_4]', 
                        'Получить системные настройки'],
            ),
	    ]
        answers = inquirer.prompt(questions)
        if answers['action'] == 'Прошить IAP и FW [MDIO_88]':
            self.write_iap_fw('bin_dio/fw.bin', 'bin_iap/iap.bin')
        elif answers['action'] == 'Прошить IAP и FW [MAI_12]':
            self.write_iap_fw('bin_ai/fw.bin', 'bin_iap/iap.bin')
        elif answers['action'] == 'Прошить IAP и FW [MAO_4]':
            self.write_iap_fw('bin_ao/fw.bin', 'bin_iap/iap.bin')

    
    def write_iap_fw(self, path_to_fw, path_to_iap):
        print(Fore.YELLOW + "Переведите устройство в режим загрузчика и нажмите 'Enter'")
        input()
        print('Начало обновления...')
        start_time = time.time()

        d = FlashAT32(self.debug_port, 115200 * 8)
        d.DEBUG_PRINT = False
        try:
            d.connect()
        except Exception as e:
            print(e)
            print(Fore.RED + 'Что-то пошло не так!!!')
            return

        print(f'Flash size: {d.get_flash_size()} kB')
        print(f'Chip ID: {d.get_uid_str()}')

        d.erase_flash()
        iap_path = path_to_iap
        fw_path = path_to_fw
        iap_path_r = 'artery_iap.bin'
        fw_path_r = 'artery_fw.bin'
        print('IPA loading...')
        d.write_file_to_flash(0x08000000, iap_path)
        print('FW loading...')
        d.write_file_to_flash(0x08021000, fw_path)

        print(Fore.GREEN + f'Запись прошла успешно за {time.time() - start_time} сек.')


class DUT:
    def __init__(self):
        pass


def main():
    tester = Tester('COM53')
    tester.menu()



if __name__ == '__main__':
    main()