from imgui_bundle import imgui, immapp, hello_imgui
import hid
import struct
import time
import enum
import subprocess
import os
import threading
import re
from tkinter import filedialog

FULL_CLEAR_BOOTLOADER_VERSION = 0xF017
# NEW_BOOTLOADER_VERSION = 0x0016

class DFUstate(enum.Enum):
    Idle = enum.auto()
    Flashing_Clear = enum.auto()
    Clear_Finished_Swapping_Bootloader = enum.auto()
    Swap_To_Clear_Failed = enum.auto()
    Clear_Failed = enum.auto()
    # Flashing_New_Bootloader = enum.auto()
    # Bootloader_Finished_Swapping_Bootloader = enum.auto()
    # Swap_To_New_Boot_Failed = enum.auto()
    # Bootloader_Failed = enum.auto()
    Flashing_New_Firmware = enum.auto()
    Firmware_Finished = enum.auto()
    Firmware_Failed = enum.auto()


def fpga_reconfig():
    try:
        bd = hid.device()
        bd.open(0x35bd,0x0101)
        bd.send_feature_report(b'\x00eB')
        bd.close()
    except OSError:
        imgui.open_popup("hid failed")

def fpga_reset():
    try:
        bd = hid.device()
        bd.open(0x35bd,0x0101)
        bd.send_feature_report(b'\x00eR')
        bd.close()
    except OSError:
        imgui.open_popup("hid failed")

def text_box(fmt: str, wrap_width = 400):
    draw_list = imgui.get_window_draw_list()
    imgui.set_cursor_pos_x(5)
    imgui.push_text_wrap_pos(imgui.get_cursor_pos_x() + wrap_width)
    imgui.text(fmt)
    imgui.pop_text_wrap_pos()
    minpos = imgui.get_item_rect_min()-imgui.ImVec2(5,5)
    maxpos = imgui.get_item_rect_max()+imgui.ImVec2(5,5)
    
    draw_list.add_rect(minpos, maxpos, imgui.IM_COL32(255,255,0,255))

class dfu_check:
    def __init__(self):
        self.dfu_util_path = None
        self.mythread = None
        self.return_status = ""
        self.success = False

    def start(self, dfu_util_path):
        self.dfu_util_path = dfu_util_path
        self.mythread = threading.Thread(target=self.run)
        self.mythread.start()

    def is_running(self):
        if self.mythread and self.mythread.is_alive():
            return True
        return False

    def run(self):
        if not self.dfu_util_path:
            print("No dfu-util path given")
            return
        full_path_to_dfu_util = os.path.realpath(os.path.join(self.dfu_util_path, "dfu-util.exe"))
        procresult = subprocess.run([
            full_path_to_dfu_util,
            "-l"],
            shell=False,
            capture_output=True
        )
        if len(procresult.stdout) > 0:
            dfu_obj_regex = re.compile(r'\[([0-9A-Fa-f]{4}):([0-9A-Fa-f]{4})\]\s*ver=([0-9A-Fa-f]{4})'.encode())
            # print(f"OUTPUT: {procresult.stdout}")
            for line in procresult.stdout.split(b'\n'):
                if line.startswith(b'Found DFU: '):
                    m = dfu_obj_regex.search(line)
                    if m:
                        vid = m.group(1).decode()
                        pid = m.group(2).decode()
                        bcdUSB = m.group(3).decode()
                        if int(vid,16) == 0x35bd and int(pid,16) == 0x0282:
                            self.dfu_version = bcdUSB
                            self.return_status = f"found dfu device, version: {bcdUSB}"
                            self.success = True
                        # print(f"Found dfu. VID:{m.group(1)}, PID:{m.group(2)}, bcdUSB:{m.group(3)}")
            if not self.success:
                self.return_status = "did not find dfu device"
        if len(procresult.stderr) > 0:
            self.return_status = f"dfu-util returned error: {procresult.stderr}"
            self.success = False

        
class dfu_load:
    def __init__(self, firmwaretype: str):
        self.firmwaretype = firmwaretype
        self.dfu_util_path = None
        self.firmware_file = None
        self.mythread = None
        self.return_status = ""
        self.success = False
        self.proc = None
        self.pct_done = None

    def start(self, dfu_util_path, firmware_file_path):
        self.dfu_util_path = dfu_util_path
        self.firmware_file = firmware_file_path
        self.mythread = threading.Thread(target=self.run)
        self.mythread.start()

    def is_running(self):
        if self.mythread and self.mythread.is_alive():
            return True
        return False
    
    def abort(self):
        if self.mythread and self.mythread.is_alive():
            if self.proc and (self.proc.poll() is None):
                self.proc.kill()

    def run(self):
        if not self.dfu_util_path:
            print("No dfu-util path given")
            return
        if not self.firmware_file:
            print("No firmware file path given")
            return
        full_path_to_dfu_util = os.path.realpath(os.path.join(self.dfu_util_path, "dfu-util.exe"))
        self.proc = subprocess.Popen(
            [full_path_to_dfu_util,
            "-D",
            self.firmware_file],
            shell = False,
            stdin=subprocess.PIPE,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT
            )
        
        upload_pct = re.compile(r"(\d+)%")
        # upload_pct = re.compile(r"(\s*(\d+)%\s*)")

        all_stdout = ""

        while self.proc.poll() is None:
            new_stdout = self.proc.stdout.read(1)
            if new_stdout:
                all_stdout += new_stdout.decode()
                m = list(upload_pct.finditer(all_stdout))
                if m:
                    self.pct_done = int(m[-1].group(1))
                # print(new_stdout)
        new_stdout = self.proc.stdout.read()
        all_stdout += new_stdout.decode()
        print(all_stdout)

        if "Download done." in all_stdout and "Done!" in all_stdout:
            self.success = True
            self.return_status = f"Finished {self.firmwaretype} download"
        else:
            self.success = False
            self.return_status = f"Error with {self.firmwaretype} download. Please retry."

class MyGui:
    def __init__(self):

        self.dfu_util_path = os.path.realpath(os.path.join(os.curdir, "dfu-util-0.9-win64"))
        self.clear_firmware_path = os.path.realpath(os.path.join(os.curdir, "full_clear_bootloader_v0.1.7.bin"))
        # self.new_bootloader_path = os.path.realpath(os.path.join(os.curdir, "bs2_fpga_bootloader_v0.1.6.bin"))
        self.new_firmware_path = os.path.realpath(os.path.join(os.curdir, "bs2-eyetracking-v0.4.9-boot-v0.1.7.bin"))

        self.full_clear_loader = None
        self.new_bootloader_loader = None
        self.new_firmware_loader = None
        self.wait_timer = None

        self.go_idle()

    def go_idle(self):
        self.dfu_state = DFUstate.Idle
        self.status_text = "Idle"
        self.dfu_checked = False
        self.dfu_checker = None

        if self.full_clear_loader and self.full_clear_loader.is_running():
            self.full_clear_loader.abort()
        self.full_clear_loader = None

        if self.new_bootloader_loader and self.new_bootloader_loader.is_running():
            self.new_bootloader_loader.abort()
        self.new_bootloader_loader = None

        if self.new_firmware_loader and self.new_firmware_loader.is_running():
            self.new_firmware_loader.abort()
        self.new_firmware_loader = None

        if self.wait_timer and self.wait_timer.is_alive():
            self.wait_timer.join()
        self.wait_timer = None

        self.num_reconfig_swaps = 0

    def gui(self):
        window_center = imgui.get_main_viewport().get_center()

        imgui.set_next_window_pos(window_center,imgui.Cond_.appearing, imgui.ImVec2(0.5,0.5))
        popup_open, _ = imgui.begin_popup_modal("hid failed")
        if popup_open:
            imgui.text("HID device open failed")
            if imgui.button("OK"):
                imgui.close_current_popup()
            imgui.end_popup()

        imgui.separator_text("File paths")
        imgui.text("dfu-util path:")
        text_box(os.path.realpath(os.path.join(self.dfu_util_path, "dfu-util.exe")))
        imgui.same_line()
        if imgui.button("...##dfu-util"):
            newpath = filedialog.askopenfilename(initialdir=self.dfu_util_path, title="Select dfu-util.exe", 
                                                 filetypes=(("Executable","*.exe"),))
            self.dfu_util_path, _ = os.path.split(newpath)

        imgui.text("full clear bootloader path:")
        text_box(self.clear_firmware_path)
        imgui.same_line()
        if imgui.button("...##clear_boot"):
            initdir, _ = os.path.split(self.clear_firmware_path)
            self.clear_firmware_path = filedialog.askopenfilename(initialdir=initdir, title="Select full_clear_bootloader.bin", 
                                                                  filetypes=(("FPGA Binary","*.bin"),))

        # imgui.text("new bootloader path:")
        # text_box(self.new_bootloader_path)
        # imgui.same_line()
        # if imgui.button("...##new_boot"):
        #     initdir, _ = os.path.split(self.new_bootloader_path)
        #     self.new_bootloader_path = filedialog.askopenfilename(initialdir=initdir, title="Select new bootloader .bin", 
        #                                                           filetypes=(("FPGA Binary","*.bin"),))

        imgui.text("new full firmware path:")
        text_box(self.new_firmware_path)
        imgui.same_line()
        if imgui.button("...##new_firm"):
            initdir, _ = os.path.split(self.new_firmware_path)
            self.new_firmware_path = filedialog.askopenfilename(initialdir=initdir, title="Select new full firmware .bin", 
                                                                  filetypes=(("FPGA Binary","*.bin"),))

        imgui.separator_text("Controls")

        if self.dfu_state != DFUstate.Idle:
            imgui.begin_disabled(True)
            overwrite_disabled = True
        else:
            overwrite_disabled = False
        if imgui.button("Start overwrite"):
            self.dfu_checked = False
            self.dfu_state = DFUstate.Flashing_Clear
        if overwrite_disabled:
            imgui.end_disabled()

        match self.dfu_state:
            case DFUstate.Idle:
                pass

            case DFUstate.Flashing_Clear:
                # sub-case - check for dfu
                if not self.dfu_checked:
                    if not self.dfu_checker:
                        self.dfu_checker = dfu_check()
                        self.dfu_checker.start(self.dfu_util_path)

                    if(self.dfu_checker.is_running()):
                        pass
                    else:
                        self.status_text = self.dfu_checker.return_status
                        if self.dfu_checker.success:
                            self.dfu_checked = True
                        else:
                            self.dfu_state = DFUstate.Clear_Failed

                else: # self.dfu_checked is True
                    if self.wait_timer:
                        self.status_text = "Waiting for reset (4 sec)..."
                        if not self.wait_timer.is_alive():
                            self.wait_timer = None
                            if self.full_clear_loader and self.full_clear_loader.is_running():
                                self.full_clear_loader.abort()
                            self.full_clear_loader = None
                            self.dfu_checker = None
                            self.dfu_checked = False
                    else:
                        if not self.full_clear_loader:
                            self.full_clear_loader = dfu_load("full clear")
                            self.full_clear_loader.start(self.dfu_util_path, self.clear_firmware_path)
                        if(self.full_clear_loader.is_running()):
                            if self.full_clear_loader.pct_done:
                                self.status_text = self.dfu_checker.return_status + "\n" + f"Loading: {self.full_clear_loader.pct_done}% done"
                        else:
                            self.status_text = self.dfu_checker.return_status + "\n" + self.full_clear_loader.return_status
                            if self.full_clear_loader.success:
                                self.dfu_checked = False
                                self.dfu_checker = None
                                self.dfu_state = DFUstate.Clear_Finished_Swapping_Bootloader
                                self.num_reconfig_swaps = 0
                            else:
                                self.dfu_state = DFUstate.Clear_Failed
                        if imgui.button("Stuck? Retry"):
                            fpga_reset()
                            self.wait_timer = threading.Thread(target=time.sleep, args=(4,))
                            self.wait_timer.start()



            case DFUstate.Clear_Failed:
                if imgui.button("Retry"):
                    if self.full_clear_loader and self.full_clear_loader.is_running():
                            self.full_clear_loader.abort()
                    self.full_clear_loader = None
                    self.dfu_checker = None
                    self.dfu_checked = False
                    self.dfu_state = DFUstate.Flashing_Clear

            case DFUstate.Clear_Finished_Swapping_Bootloader:
                # first see what DFU is loaded
                if not self.dfu_checked: 
                    if not self.dfu_checker:
                        # create the checker task
                        self.dfu_checker = dfu_check()
                        self.dfu_checker.start(self.dfu_util_path)
                    if(self.dfu_checker.is_running()):
                        pass
                    else:
                        self.status_text = self.dfu_checker.return_status
                        if self.dfu_checker.success:
                            self.dfu_checked = True
                        else:
                            self.dfu_state = DFUstate.Swap_To_Clear_Failed
                else:

                    if self.wait_timer:
                        self.status_text = "Waiting for reconfig (4 sec)..."
                        if not self.wait_timer.is_alive():
                            self.wait_timer = None
                            self.dfu_checker = None
                            self.dfu_checked = False

                    elif int(self.dfu_checker.dfu_version, 16) != FULL_CLEAR_BOOTLOADER_VERSION:
                        # wrong DFU loaded, try reconfig and go again
                        if not self.wait_timer:
                            if self.num_reconfig_swaps > 3:
                                # couldn't find the special bootloader, consider this a failure
                                self.dfu_checker = None
                                self.dfu_checked = False
                                self.dfu_state = DFUstate.Clear_Failed
                            else:
                                fpga_reconfig()
                                self.num_reconfig_swaps += 1
                                self.wait_timer = threading.Thread(target=time.sleep, args=(4,))
                                self.wait_timer.start()
                    else:
                        # correct DFU! continue to loading the new firmware
                        self.dfu_state = DFUstate.Flashing_New_Firmware
                        self.dfu_checker = None
                        self.dfu_checked = False

            case DFUstate.Swap_To_Clear_Failed:
                if imgui.button("Retry"):
                    if self.full_clear_loader and self.full_clear_loader.is_running():
                            self.full_clear_loader.abort()
                    self.full_clear_loader = None
                    self.dfu_checker = None
                    self.dfu_checked = False
                    self.dfu_state = DFUstate.Clear_Finished_Swapping_Bootloader
        
            # case DFUstate.Flashing_New_Bootloader:
            #     # sub-case - check for dfu
            #     if not self.dfu_checked:
            #         if not self.dfu_checker:
            #             self.dfu_checker = dfu_check()
            #             self.dfu_checker.start(self.dfu_util_path)

            #         if(self.dfu_checker.is_running()):
            #             pass
            #         else:
            #             self.status_text = self.dfu_checker.return_status
            #             if self.dfu_checker.success:
            #                 self.dfu_checked = True
            #             else:
            #                 self.dfu_state = DFUstate.Bootloader_Failed
            #     else: # self.dfu_checked is True
            #         if self.wait_timer:
            #             self.status_text = "Waiting for reset (4 sec)..."
            #             if not self.wait_timer.is_alive():
            #                 self.wait_timer = None
            #                 if self.new_bootloader_loader and self.new_bootloader_loader.is_running():
            #                     self.new_bootloader_loader.abort()
            #                 self.new_bootloader_loader = None
            #                 self.dfu_checker = None
            #                 self.dfu_checked = False
            #         else:
            #             if not self.new_bootloader_loader:
            #                 self.new_bootloader_loader = dfu_load("new bootloader")
            #                 self.new_bootloader_loader.start(self.dfu_util_path, self.new_bootloader_path)
            #             if(self.new_bootloader_loader.is_running()):
            #                 if self.new_bootloader_loader.pct_done:
            #                     self.status_text = self.dfu_checker.return_status + "\n" + f"Loading: {self.new_bootloader_loader.pct_done}% done"
            #             else:
            #                 self.status_text = self.dfu_checker.return_status + "\n" + self.new_bootloader_loader.return_status
            #                 if self.new_bootloader_loader.success:
            #                     self.dfu_checked = False
            #                     self.dfu_checker = None
            #                     self.dfu_state = DFUstate.Bootloader_Finished_Swapping_Bootloader
            #                     self.num_reconfig_swaps = 0
            #                 else:
            #                     self.dfu_state = DFUstate.Bootloader_Failed
            #             if imgui.button("Stuck? Retry"):
            #                 fpga_reset()
            #                 self.wait_timer = threading.Thread(target=time.sleep, args=(4,))
            #                 self.wait_timer.start()

            # case DFUstate.Bootloader_Failed:
            #     if imgui.button("Retry"):
            #         if self.new_bootloader_loader and self.new_bootloader_loader.is_running():
            #                 self.new_bootloader_loader.abort()
            #         self.new_bootloader_loader = None
            #         self.dfu_checker = None
            #         self.dfu_checked = False
            #         self.dfu_state = DFUstate.Flashing_New_Bootloader

            # case DFUstate.Bootloader_Finished_Swapping_Bootloader:
            #     # first see what DFU is loaded
            #     if not self.dfu_checked: 
            #         if not self.dfu_checker:
            #             # create the checker task
            #             self.dfu_checker = dfu_check()
            #             self.dfu_checker.start(self.dfu_util_path)
            #         if(self.dfu_checker.is_running()):
            #             pass
            #         else:
            #             self.status_text = self.dfu_checker.return_status
            #             if self.dfu_checker.success:
            #                 self.dfu_checked = True
            #             else:
            #                 self.dfu_state = DFUstate.Swap_To_New_Boot_Failed
            #     else:
            #         if self.wait_timer:
            #             self.status_text = "Waiting for reconfig (4 sec)..."
            #             if not self.wait_timer.is_alive():
            #                 self.wait_timer = None
            #                 self.dfu_checker = None
            #                 self.dfu_checked = False

            #         elif int(self.dfu_checker.dfu_version, 16) != NEW_BOOTLOADER_VERSION:
            #             # wrong DFU loaded, try reconfig and go again
            #             if not self.wait_timer:
            #                 if self.num_reconfig_swaps > 3:
            #                     # couldn't find the special bootloader, consider this a failure
            #                     self.dfu_checker = None
            #                     self.dfu_checked = False
            #                     self.dfu_state = DFUstate.Bootloader_Failed
            #                 else:
            #                     fpga_reconfig()
            #                     self.num_reconfig_swaps += 1
            #                     self.wait_timer = threading.Thread(target=time.sleep, args=(4),)
            #                     self.wait_timer.start()
            #         else:
            #             # correct DFU! continue to loading the new bootloader
            #             self.dfu_state = DFUstate.Flashing_New_Firmware
            #             self.dfu_checker = None
            #             self.dfu_checked = False

            # case DFUstate.Swap_To_New_Boot_Failed:
            #     if imgui.button("Retry"):
            #         if self.full_clear_loader and self.full_clear_loader.is_running():
            #                 self.full_clear_loader.abort()
            #         self.full_clear_loader = None
            #         self.dfu_checker = None
            #         self.dfu_checked = False
            #         self.dfu_state = DFUstate.Bootloader_Finished_Swapping_Bootloader
        

            case DFUstate.Flashing_New_Firmware:
                # sub-case - check for dfu
                if not self.dfu_checked:
                    if not self.dfu_checker:
                        self.dfu_checker = dfu_check()
                        self.dfu_checker.start(self.dfu_util_path)

                    if(self.dfu_checker.is_running()):
                        pass
                    else:
                        self.status_text = self.dfu_checker.return_status
                        if self.dfu_checker.success:
                            self.dfu_checked = True
                        else:
                            self.dfu_state = DFUstate.Firmware_Failed
                else: # self.dfu_checked is True
                    if self.wait_timer:
                        self.status_text = "Waiting for reset (4 sec)..."
                        if not self.wait_timer.is_alive():
                            self.wait_timer = None
                            if self.new_firmware_loader and self.new_firmware_loader.is_running():
                                self.new_firmware_loader.abort()
                            self.new_firmware_loader = None
                            self.dfu_checker = None
                            self.dfu_checked = False
                    else:
                        if not self.new_firmware_loader:
                            self.new_firmware_loader = dfu_load("new firmware")
                            self.new_firmware_loader.start(self.dfu_util_path, self.new_firmware_path)
                        if(self.new_firmware_loader.is_running()):
                            if self.new_firmware_loader.pct_done:
                                self.status_text = self.dfu_checker.return_status + "\n" + f"Loading: {self.new_firmware_loader.pct_done}% done"
                        else:
                            self.status_text = self.dfu_checker.return_status + "\n" + self.new_firmware_loader.return_status
                            if self.new_firmware_loader.success:
                                self.dfu_checked = False
                                self.dfu_checker = None
                                self.dfu_state = DFUstate.Firmware_Finished
                                self.num_reconfig_swaps = 0
                            else:
                                self.dfu_state = DFUstate.Firmware_Failed
                        if imgui.button("Stuck? Retry"):
                            fpga_reset()
                            self.wait_timer = threading.Thread(target=time.sleep, args=(4,))
                            self.wait_timer.start()

            case DFUstate.Firmware_Failed:
                if imgui.button("Retry"):
                    if self.new_firmware_loader and self.new_firmware_loader.is_running():
                            self.new_firmware_loader.abort()
                    self.new_firmware_loader = None
                    self.dfu_checker = None
                    self.dfu_checked = False
                    self.dfu_state = DFUstate.Flashing_New_Firmware

            case DFUstate.Firmware_Finished:
                imgui.text_colored(imgui.ImVec4(0.0, 1.0, 0.0, 1.0), "Complete!")
                if imgui.button("Reload for next unit"):
                    self.go_idle()

        imgui.text("STATUS:")
        imgui.text_wrapped(self.status_text)

        imgui.separator()
        imgui.separator_text("Extra Controls")
        imgui.push_style_color(imgui.Col_.text, imgui.IM_COL32(255,100,255,255))
        imgui.text_wrapped("WARNING: Abort or reconfig after new firmware is partially written can result in a bricked device. "
                           "JTAG connection will be necessary to fix. "
                           "Try the \"Stuck? Retry\" button first.")
        imgui.pop_style_color()
        if imgui.button("Abort. Fully reset to the beginning!"):
            self.go_idle()
        if self.dfu_state != DFUstate.Idle:
            imgui.begin_disabled(True)
        imgui.same_line()
        if imgui.button("Toggle Reconfig"):
            fpga_reconfig()
        if self.dfu_state != DFUstate.Idle:
            imgui.end_disabled()

mygui = MyGui()

immapp.run(
    gui_function=mygui.gui,  # The Gui function to run
    window_title="update FPGA bootloader",  # the window title
    # window_size_auto=True,  # Auto size the application window given its widgets
    window_size=(600,500)
    # Uncomment the next line to restore window position and size from previous run
    # window_restore_previous_geometry==True
)