import os, sys

if getattr(sys, "frozen", False):
    exe_dir = os.path.dirname(sys.executable)
    os.environ["PATH"] = exe_dir + os.pathsep + os.environ.get("PATH", "")

import os, sys, ctypes, subprocess
import threading
import builtins

def is_admin():
    try:
        return ctypes.windll.shell32.IsUserAnAdmin()
    except:
        return False
    
if not is_admin():
    ctypes.windll.shell32.ShellExecuteW(
        None, u"runas", str(sys.executable), str(__file__), None, 1
    )
    sys.exit()

if os.path.exists("./lighthouse_console.exe"):
    setattr(builtins, 'valve_tools_path', './')
    setattr(builtins, 'steam_vr_path', './')

import time
import logging
import queue
import re
from gui_log import *
from imgui_bundle import imgui, immapp, hello_imgui
from plyer import notification
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass
from random import Random
from typing import Any, Optional
from enum import Enum, auto
import hid
from steamvr.unbuffered.DeviceMonitor import instance

import bs_usb_tools
import bs_hmd_tools
from quick_updater_settings import config

logging.getLogger()

#tracking_dir = r'C:\Users\facto\calibot\venv\Lib\site-packages\steamvr\tools\bin\win64'
#init_json = r"C:\Users\facto\calibot\venv\Lib\site-packages\steamvr\init_json_evt.json"
init_json = r".\init_json_config.json"

QUICK_UPDATER_VERSION = '1.9'

RAND = Random()
VIDS_ATMEL = [0x35bd]
PIDS_ATMEL = [0x0101, 0x4004]  # , 0x4004
VIDS_TUNDRA = [0x28de]
PIDS_TUNDRA = [0x2300]
KNOWN_VIDS = VIDS_ATMEL + VIDS_TUNDRA
KNOWN_PIDS = PIDS_ATMEL + PIDS_TUNDRA

# Values derived from the yaml config
firmware_path = config.firmware_path
fpga_firmware_path = config.fpga_firmware_path
CLEAR_HISTORY = config.clear_device_monitor_history
USB_CLEANUP_THRESHOLD = config.usb_cleanup_threshold
MAX_THREADS = config.max_threads
DELAY_MCU = config.delays.mcu_seconds
DELAY_LH = config.delays.lighthouse_seconds
UPDATE_INTERVAL = config.delays.firmware_update_interval

MAX_DETECTIONS_PER_DEVICE = 10

window_title = ""
success_count = 0
fail_count = 0
last_clean_device_qty = 0
update_timeslot_current = 0
tick = 0
tick_lock = threading.Lock()
update_timestamps = []
update_timestamps_lock = threading.Lock()
target_firmware_path = ""
target_firmware_version = ""
firmware_payload = None
target_fpga_firmware_path = ""
target_fpga_firmware_version = ""
fpga_firmware_payload = None
BOOTLOADER_FPGA_FW_VERSION = "0.1.7"

icon_path_black = os.path.join(sys._MEIPASS, 'bynd-black.ico') if hasattr(sys, '_MEIPASS') else 'bynd-black.ico'
icon_path_purple = os.path.join(sys._MEIPASS, 'bynd-purple.ico') if hasattr(sys, '_MEIPASS') else 'bynd-purple.ico'
icon_path_orange = os.path.join(sys._MEIPASS, 'bynd-orange.ico') if hasattr(sys, '_MEIPASS') else 'bynd-orange.ico'
icon_path_clear = os.path.join(sys._MEIPASS, 'bynd-clear.ico') if hasattr(sys, '_MEIPASS') else 'bynd-clear.ico'
_HMD_SERIAL_PATTERNS = {
    re.compile(r'^[Bb][Ss]2E?B', re.IGNORECASE): icon_path_black,
    re.compile(r'^[Bb][Ss]2E?P', re.IGNORECASE): icon_path_purple,
    re.compile(r'^[Bb][Ss]2E?O', re.IGNORECASE): icon_path_orange,
}

# Queue for processing HID devices asynchronously (to avoid blocking WMI watcher thread)
device_queue = queue.Queue()


class LEDModes(Enum):
    DEFAULT = 0
    SOLID = 1
    BREATHING = 2
    BLINKING = 3


class UpdateState(Enum):
    ONGOING = auto()
    FPGA = auto()
    BOOTLOADER = auto()
    ERROR = auto()
    DONE = auto()
quiet_update_states = [UpdateState.ERROR, UpdateState.DONE]


def update_window_title():
    global window_title
    window_title = f"BEYOND QUICK UPDATER TOOL ::: v{QUICK_UPDATER_VERSION}"

update_window_title()
    

def get_target_firmware_path():
    global firmware_path
    fw_path = firmware_path
    fw_file_ver = ""
    log_text = ""

    # search working dir for override file
    for f in os.listdir('.'):
        if os.path.isfile(f) and 'override' in f.lower() and f.lower().endswith('.beyondfw'):
            fw_path = os.path.abspath(f)
            log_text = f"Using override firmware file from working directory"
            break
    if log_text == "":
        if fw_path != "":
            log_text = f"Using configured firmware"
        else:
            utility_path = bs_hmd_tools.get_utility_beyondfw_path()
            if utility_path and os.path.exists(utility_path):
                fw_path = utility_path
                log_text = f"Using utility firmware"

    assert os.path.isfile(fw_path), f"Firmware file not found: {fw_path} ({log_text})"

    fw_file_ver = bs_hmd_tools.get_file_fw_version(fw_path)
    log_debug(log_text + f" (v{fw_file_ver}): {fw_path}")
    return fw_path
    

def get_target_fpga_firmware_path():
    global fpga_firmware_path
    fw_path = fpga_firmware_path
    fw_file_ver = ""
    log_text = ""

    # search working dir for override file
    for f in os.listdir('.'):
        if os.path.isfile(f) and 'override' in f.lower() and f.lower().endswith('.dfu'):
            fw_path = os.path.abspath(f)
            log_text = f"Using override FPGA firmware file from working directory"
            break
    if log_text == "":
        if fw_path != "":
            log_text = f"Using configured FPGA firmware"
        else:
            utility_path = bs_hmd_tools.get_utility_etfw_path()
            if utility_path and os.path.exists(utility_path):
                fw_path = utility_path
                log_text = f"Using utility FPGA firmware"

    assert os.path.isfile(fw_path), f"FPGA firmware file not found: {fw_path} ({log_text})"

    fw_file_ver = bs_hmd_tools.get_dfu_file_fw_version(fw_path)
    log_debug(log_text + f" (v{fw_file_ver}): {fw_path}")
    return fw_path


def icon_path_from_sn(hmd_serial: str) -> str:
    """
    Determine the icon path based on the HMD serial number color code.
    
    Serial number format: BS##X... where:
        - BS: Fixed prefix (case-insensitive)
        - ##: Two digits
        - X: Color code - B(lack), P(urple), O(range)
    
    Args:
        hmd_serial: The HMD serial number string
        
    Returns:
        Path to the appropriate icon file based on color code,
        or the clear icon path if no match is found or input is invalid
    """
    if not hmd_serial or not isinstance(hmd_serial, str):
        return icon_path_clear
    
    for pattern, icon_path in _HMD_SERIAL_PATTERNS.items():
        if pattern.match(hmd_serial):
            return icon_path
    
    return icon_path_clear


def config_reload():
    global firmware_path, fpga_firmware_path, CLEAR_HISTORY, USB_CLEANUP_THRESHOLD, MAX_THREADS
    global DELAY_MCU, DELAY_LH, UPDATE_INTERVAL
    global window_title

    if config.reload():
        firmware_path = config.firmware_path
        fpga_firmware_path = config.fpga_firmware_path
        CLEAR_HISTORY = config.clear_device_monitor_history
        USB_CLEANUP_THRESHOLD = config.usb_cleanup_threshold
        MAX_THREADS = config.max_threads
        DELAY_MCU = config.delays.mcu_seconds
        DELAY_LH = config.delays.lighthouse_seconds
        UPDATE_INTERVAL = config.delays.firmware_update_interval
        window_title = f"BEYOND QUICK UPDATER TOOL v{QUICK_UPDATER_VERSION}"
        log_info(f"Configuration reloaded")


def static_vars(**kwargs):
    def decorate(func):
        for k in kwargs:
            setattr(func, k, kwargs[k])
        return func
    return decorate


@dataclass
class HeadsetHID:
    atmel: Any = None
    lh_serial: str = ''
    hmd_serial: str = ''
    tundra_path: str = ''
    bootloader_path: str = ''
    fw_version: str = ''
    fpga_fw_version: str = ''
    update_state: UpdateState = UpdateState.ONGOING
    detection_counter: int = 0
    unresolved_hid_counter: int = 0
    temp_hid_device: Optional[hid.device] = None

    def __repr__(self):
        return f"{self.atmel['serial_number']} - {self.lh_serial}"
    
    def safety_check(self):
        assert self.atmel is not None
        assert self.lh_serial != ''
        assert 'port_chain' in self.atmel

    def get_dynamic_hid_path(self):
        if self.bootloader_path != "":
            return self.bootloader_path
        else:
            return self.atmel['path']

    def get_hid_device(self):
        retry_time = DELAY_MCU
        target_device_path = self.get_dynamic_hid_path()
        if self.temp_hid_device is None:
            self.temp_hid_device = hid.device()
        else:
            self.temp_hid_device.close()
        while retry_time > 0:
            try:
                self.temp_hid_device.open_path(target_device_path)
                return self.temp_hid_device
            except OSError as e:
                retry_time -= 0.5
                time.sleep(0.5)
        return self.temp_hid_device
    
    def clear_hid_device(self):
        if self.temp_hid_device is not None:
            self.temp_hid_device.close()

    def send_feature_report(self, data: bytes):
        retries = 3
        while retries > 0:
            try:
                device = self.get_hid_device()
                result = device.send_feature_report(data)
                self.clear_hid_device()
                return result
            except OSError as e:
                log_warning(f"{self.hmd_serial}: OSError sending HID feature report: {e}. Retrying...")
                retries -= 1
                time.sleep(0.5)
            finally:
                self.clear_hid_device()

    def set_light_color(self, rgb: tuple[int, int, int], led_mode: LEDModes = LEDModes.DEFAULT):
        target_device_path = self.get_dynamic_hid_path()
        assert bs_hmd_tools.device_is_beyond(target_device_path) == 1, "Cannot set light color on given device/in its current state"
        self.send_feature_report(bytes([0, ord('L'), rgb[0], rgb[1], rgb[2], led_mode.value]))

    def refresh_hmd_serial(self):
        """ Fetches and updates the hmd_serial attribute. """
        self.hmd_serial = bs_hmd_tools.get_hmd_serial(direct_device=self.get_hid_device())
        self.clear_hid_device()
        return self.hmd_serial

    def refresh_fw_version(self):
        """ Fetches and updates the fw_version attribute. Waits up to DELAY_MCU seconds for the device to be ready """
        self.fw_version = bs_hmd_tools.get_hmd_fw_version(direct_device=self.get_hid_device(), retry_time=DELAY_MCU)
        self.clear_hid_device()
        return self.fw_version

    def refresh_fpga_fw_version(self):
        """ Fetches and updates the fpga_fw_version attribute. """
        self.fpga_fw_version = bs_hmd_tools.get_hmd_fpga_fw_version(direct_device=self.get_hid_device())
        self.clear_hid_device()
        return self.fpga_fw_version
    
    def enter_hmd_bootloader(self, confirm_bootloader: bool = True) -> Optional[bytes]:
        """ Sends command to enter bootloader mode. Returns new bootloader path if successful. """
        bootloader_path = bs_hmd_tools.enter_hmd_bootloader(direct_device=self.get_hid_device(), confirm_bootloader=confirm_bootloader)
        self.clear_hid_device()
        return bootloader_path
    
    def update_hmd_firmware(self, firmware_payload: bytes) -> bool:
        """ Updates the HMD firmware using the provided payload. Returns True if successful. """
        self.refresh_fw_version()  # wait for device to be ready
        result = bs_hmd_tools.update_hmd_firmware(firmware_payload, bootloader_path=self.get_dynamic_hid_path())
        return result
    
    def restart_hub(self):
        self.safety_check()
        log_info(f"Restarting USB hub for device {self.atmel['serial_number']} - {self.lh_serial}")
        print(f"Restarting USB hub for device {self.atmel['serial_number']} - {self.lh_serial}")
        parent_hub_chain = '-'.join(self.atmel['port_chain'][:-2])
        bs_usb_tools.restart_port(parent_hub_chain)

    def erase_records(self, include_headsets_log = False):
        global matching
        matching.forget_headset(headset=self, include_headsets_log=include_headsets_log)
    
    def reenqueue(self, bootloader_path: Optional[str] = None):
        # self.safety_check()
        global matching
        global update_timeslot_current

        if bootloader_path is not None:
            self.bootloader_path = bootloader_path

        if update_timeslot_current < tick:
            update_timeslot_current = tick - UPDATE_INTERVAL
        update_timeslot_current += UPDATE_INTERVAL
        matching.on_match(self, update_timeslot_current)


class HeadsetHIDMatching:
    def __init__(self, on_match):
        self.prospect = None
        self.atmel_candidates = []
        self.tundra_candidates = []
        self.matched_paths = []
        self.processed_headsets = []
        self.on_match = on_match
        self._lock = threading.Lock()

    def __on_match(self, dev):
        global tick
        global update_timeslot_current

        if update_timeslot_current < tick:
            update_timeslot_current = tick - UPDATE_INTERVAL
        update_timeslot_current += UPDATE_INTERVAL

        if dev.atmel:
            # log_debug(f"Found match for HMD {dev.atmel['serial_number']} - {dev.lh_serial}")
            print(f"Found match for HMD {dev.atmel['serial_number']} - {dev.lh_serial}")
            # log_info(f"[TICK:{tick}] Scheduling update for {dev.lh_serial} at tick {update_timeslot_current}.")  # debug
        self.on_match(dev, update_timeslot_current)  # update_hmd_firmware() will offset this

    def fetch_port_chain(self, serial_number: str):
        """Fetches the port chain for Atmel of a given serial number using bs_usb_tools"""
        usb_info = bs_usb_tools.get_info(f"*{serial_number}*")
        retries = 3
        while 'PortChain' not in usb_info and retries > 0:
            time.sleep(1)
            usb_info = bs_usb_tools.get_info(f"*{serial_number}*")
            retries -= 1
        
        if 'PortChain' in usb_info:
            return usb_info['PortChain']
        return None
    
    def forget_headset(self, headset: HeadsetHID, include_headsets_log = False):
        """Removes a headset from processed list and matched paths"""
        with self._lock:
            if headset in self.processed_headsets and include_headsets_log:
                self.processed_headsets.remove(headset)
            if headset.atmel and headset.atmel['path'] in self.matched_paths:
                self.matched_paths.remove(headset.atmel['path'])
            if headset.bootloader_path in self.matched_paths:
                self.matched_paths.remove(headset.bootloader_path)
            if headset.tundra_path in self.matched_paths:
                self.matched_paths.remove(headset.tundra_path)
    
    def clean_matched_paths(self):
        """Removes entries for devices that are no longer connected"""
        # using hid.enumerate() to get current paths
        atmel_devs = hid.enumerate(VIDS_ATMEL[0])
        # populate tundra but filter out non-MI_00 controllers
        tundra_devs = [dev for dev in hid.enumerate(VIDS_TUNDRA[0],PIDS_TUNDRA[0]) if dev['product_string'] == 'Controller']
        current_paths = {dev['path'] for dev in atmel_devs}
        current_paths.update({dev['path'] for dev in tundra_devs})
        with self._lock:
            self.matched_paths = [path for path in self.matched_paths if path in current_paths]

    def on_new_hid_device(self, dev, i, new_devs, prev_devs):
        """Non-blocking callback - just enqueues the device for async processing"""
        vid = dev['vendor_id']
        pid = dev['product_id']
        product = dev['product_string']
        path = dev['path']
        
        # Quick filter check before enqueueing
        self.clean_matched_paths()
        if vid in KNOWN_VIDS and pid in KNOWN_PIDS and path not in self.matched_paths:
            if vid in VIDS_TUNDRA and product != 'Controller':
                # log_debug(f"Skipping non-MI_02 Tundra device: Serial={serial}, Path={path}")  # debug
                return []
            # Enqueue for async processing to avoid blocking WMI watcher thread
            device_queue.put((dev, i, new_devs, prev_devs))
        
        return []

    def process_device(self, dev, i, new_devs, prev_devs):
        """Process a device from the queue - runs in worker thread"""
        global target_firmware_version, quiet_update_states
        vid = dev['vendor_id']
        pid = dev['product_id']
        product = dev['product_string']
        serial = dev['serial_number'].strip()
        path = dev['path']
        matched_headset = None
        return_devs = []
        new_candidate = False

        if product != 'Bigscreen USB Firmware Upgrade':  # following doesn't apply to bootloader devices
            with self._lock:
                if path in self.matched_paths:
                    # Not much point to this next line idk, maybe nix later
                    matched_headset = next((h for h in self.processed_headsets if h.atmel and h.atmel['path'] == path), None)
                    return return_devs

        matched_headset = None
        with self._lock:
            if vid in VIDS_TUNDRA and pid in PIDS_TUNDRA and product == 'Controller':
                dev['port_chain'] = self.fetch_port_chain(serial)
                if dev['port_chain'] is None:
                    log_warning(f"Could not get PortChain for device with serial {serial}")
                    return return_devs
                # log_debug(f"Adding Tundra candidate: Serial={serial}")  # debug
                self.tundra_candidates.append(dev)
                new_candidate = True
            if vid in VIDS_ATMEL and pid in PIDS_ATMEL:
                if product == 'Bigscreen USB Firmware Upgrade':
                    # notification.notify(f"Wanderer detected!", f"This headset started in bootloader mode (っ´ཀ`)っ Attempting to update...", "Beyond Quick Updater", timeout=10, app_icon=icon_path_black)
                    log_info(f"Bootloader device found. Processing!")  # debug
                    matched_headset = HeadsetHID(bootloader_path=path, hmd_serial='=,,= (wanderer)')  # create husk record to update now
                    self.matched_paths.append(path)
                if product == 'Beyond':
                    if any(h.atmel and h.atmel['serial_number'] == serial and h.tundra_path != '' for h in self.processed_headsets):
                        matched_headset = next(h for h in self.processed_headsets if h.atmel and h.atmel['serial_number'] == serial and h.tundra_path != '')
                        self.matched_paths.extend([matched_headset.atmel['path'], matched_headset.tundra_path])
                        if matched_headset.update_state not in quiet_update_states:
                            log_debug(f"Existing headset found for Atmel serial {serial}, jumping straight to evaluation") # debug
                    else:
                        dev['port_chain'] = self.fetch_port_chain(serial)
                        if dev['port_chain'] is None:
                            log_warning(f"Could not get PortChain for device with serial {serial}")
                            return return_devs
                        # log_debug(f"Adding Atmel candidate: Serial={serial}")  # debug
                        self.atmel_candidates.append(dev)
                        new_candidate = True

            # check that we found at least one of each. then check for matched pairs
            if new_candidate and len(self.atmel_candidates) > 0 and len(self.tundra_candidates) > 0:
                # log_debug(f"Checking for matches")  # debug
                for a_i, a_dev in enumerate(self.atmel_candidates):
                    # log_debug(f"Atmel candidate: {a_i}")  # debug
                    for t_i, t_dev in enumerate(self.tundra_candidates):
                        # log_debug(f"  Tundra candidate: {t_i}")  # debug
                        parent_ports_match = False
                        # iterate through port chains to see if they match up to index -3 (all but the last 2)
                        if len(a_dev['port_chain']) >= 3 and len(t_dev['port_chain']) >= 3:
                            parent_ports_match = a_dev['port_chain'][:-2] == t_dev['port_chain'][:-2]
                        if not parent_ports_match:
                            # log_debug(f"    USB hubs do not match: {a_dev['port_chain']} != {t_dev['port_chain']}")  # debug
                            continue
                        if a_dev['path'] not in self.matched_paths and t_dev['path'] not in self.matched_paths:
                            log_debug(f"    Found match!")  # debug
                            print(f"With atmel: {a_dev}")
                            print(f"With tundra: {t_dev}")
                            matched_headset = HeadsetHID(a_dev, t_dev['serial_number'], tundra_path=t_dev['path'])
                            del self.atmel_candidates[a_i]
                            del self.tundra_candidates[t_i]
                            self.matched_paths.extend([a_dev['path'], t_dev['path']])
                            break  # Exit inner loop after finding a match
                        else:
                            log_debug(f"    One of the devices was already processed")
                    if matched_headset:
                        break  # Exit outer loop after finding a match

        # Call on_match OUTSIDE the lock to avoid blocking
        if matched_headset:
            self.processed_headsets.append(matched_headset)
            self.__on_match(matched_headset)

        # print(f"---\r\nPrevious devices list: {self.matched_paths}")  # debug
        # print(f"Returning devices: {return_devs}")  # debug
        return return_devs

    def reset(self):
        self.prospect = None
        self.atmel_candidates = []
        self.tundra_candidates = []
        self.matched_paths = []


class TimerState:
    def __init__(self):
        self.reset()

    def activate(self):
        self.activated = True

    def end(self):
        self.finished = True

    def reset(self):
        self.finished = False
        self.activated = False


def set_light_color(dev: HeadsetHID, rgb: tuple[int, int, int], led_mode: LEDModes = LEDModes.DEFAULT):
    hmd_path = dev.atmel['path']
    assert bs_hmd_tools.device_is_beyond(hmd_path) == 1, "Cannot set light color on given device/in its current state"
    device = hid.device()
    try:
        device.open_path(hmd_path)
        print(f'Setting led on HMD {dev.hmd_serial} to {rgb}')
        device.send_feature_report(bytes([0, ord('L'), rgb[0], rgb[1], rgb[2], led_mode.value]))
        device.close()
    except Exception as e:
        pass
    finally:
        try:
            device.close()
        except:
            pass


def update_hmd_firmware(headset: HeadsetHID, update_timeslot: int = 0):
    global tick
    global success_count, fail_count
    global update_timestamps, update_timestamps_lock
    start_tick = tick
    global quiet_update_states
    global target_firmware_version, firmware_payload
    global target_fpga_firmware_version, fpga_firmware_payload

    on_latest_fw = False
    has_et_hardware = False
    on_latest_fpga_fw = False
    is_bootloader = False

    target_device_path = headset.get_dynamic_hid_path()
    is_bootloader = bs_hmd_tools.device_is_beyond(target_device_path) == 2

    # headset.clear_hid_device()
    # time.sleep(DELAY_MCU)  # give the MCU time to post
    headset.refresh_fw_version()  # blocks for up to DELAY_MCU seconds

    if(headset.update_state not in quiet_update_states):
        headset.detection_counter += 1
    
    if not is_bootloader:
        if headset.hmd_serial == "":
            headset.refresh_hmd_serial()
        # Verify HMD serial
        if headset.hmd_serial:
            if headset.update_state not in quiet_update_states:
                log_debug(f"Found HMD: {headset.hmd_serial}")
            has_et_hardware = headset.hmd_serial.startswith('BS2E')
        else:
            log_warning(f"HMD {headset.atmel['serial_number']} serial not found or not set")
            headset.hmd_serial = "MISSING_SERIAL"

        # Sometimes HMDs get in a state where HID commands trigger a reenumeration. Catch that here.
        # if headset.unresolved_hid_counter >= 1 and headset.unresolved_hid_counter < 5:
        #     if headset.unresolved_hid_counter % 2 == 1:
        #         headset.unresolved_hid_counter += 1
        #         log_warning(headset.hmd_serial + ": Unable to resolve HID message. Restarting root hub to try updating again...")
        #         headset.erase_records()
        #         instance.reset_history()
        #         headset.restart_hub()
        # elif headset.unresolved_hid_counter >= 5:
        #     if headset.update_state not in quiet_update_states:
        #         log_error(f"{headset.hmd_serial}: HID device is stuck as claimed. This device will be updatable after a software restart.")
        #         headset.update_state = UpdateState.ERROR
        #         fail_count += 1
        #     headset.set_light_color(rgb=(255, 0, 0), led_mode=LEDModes.BLINKING)
        #     return

        # headset.unresolved_hid_counter += 1
        # headset.fw_version = bs_hmd_tools.get_hmd_fw_version(target_device_path)
        # headset.unresolved_hid_counter = 0
        on_latest_fw = headset.fw_version == target_firmware_version
        if headset.update_state not in quiet_update_states:
            log_info(f"{headset.hmd_serial}: Current FW is v{headset.fw_version}")
        if has_et_hardware:
            headset.refresh_fpga_fw_version()
            on_latest_fpga_fw = headset.fpga_fw_version == target_fpga_firmware_version

        # Catch devices that are failing updates repeatedly
        if headset.detection_counter >= MAX_DETECTIONS_PER_DEVICE and (not on_latest_fw or (has_et_hardware and not on_latest_fpga_fw)):
            if headset.update_state != UpdateState.ERROR:
                failed_firmwares = f"HMD" if not on_latest_fw else ""
                if has_et_hardware and not on_latest_fpga_fw:
                    if failed_firmwares != "":
                        failed_firmwares += " or "
                    failed_firmwares += f"ET"
                log_warning(f"{headset.hmd_serial}: Device has been detected {headset.detection_counter} times without successful updates to {failed_firmwares}. Aborting further update attempts.")
                headset.update_state = UpdateState.ERROR
                fail_count += 1
            headset.set_light_color(rgb=(255, 0, 0), led_mode=LEDModes.SOLID)
            return

        if on_latest_fw:
            if has_et_hardware and headset.fpga_fw_version != "":
                if headset.update_state not in quiet_update_states:
                    if headset.fpga_fw_version == BOOTLOADER_FPGA_FW_VERSION:
                        log_info(f"{headset.hmd_serial}: HMD FPGA is in bootloader mode")
                    else:
                        log_info(f"{headset.hmd_serial}: Current FPGA FW is v{headset.fpga_fw_version}")
                if not on_latest_fpga_fw:
                    headset.set_light_color(rgb=(255, 255, 255), led_mode=LEDModes.BLINKING)
                    headset.update_state = UpdateState.FPGA
                    log_info(headset.hmd_serial + ": " + f"Now updating FPGA firmware to v{target_fpga_firmware_version}...")
                    direct_dev = headset.get_hid_device()
                    dfu_port_chain = bs_hmd_tools.enter_fpga_dfu(direct_device=direct_dev)
                    fpga_done = bs_hmd_tools.update_fpga_firmware(fpga_firmware_payload, dfu_port_chain, reset_direct_device=direct_dev)
                    if not fpga_done:
                        headset.clear_hid_device()
                        headset.erase_records()
                        instance.reset_history()
                        headset.restart_hub()
                        return
                    headset.fpga_fw_version = bs_hmd_tools.get_hmd_fpga_fw_version(direct_device=direct_dev)
                    headset.clear_hid_device()
                    if headset.update_state not in quiet_update_states:
                        log_info(headset.hmd_serial + f": FPGA is now running firmware v{headset.fpga_fw_version}")
            
            headset.set_light_color(rgb=(0, 255, 0), led_mode=LEDModes.SOLID)
            if headset.update_state not in quiet_update_states:
                log_info(headset.hmd_serial + ": " + f"HMD now on all target firmwares")
                notification.notify(f"{headset.hmd_serial} is updated", f"♪┏(・o・)┛♪ All good~ ::: FW v{target_firmware_version} ::: FPGA v{target_fpga_firmware_version}", "Beyond Quick Updater", timeout=5, app_icon=icon_path_from_sn(headset.hmd_serial))
                headset.update_state = UpdateState.DONE
                success_count += 1
            with update_timestamps_lock:
                update_timestamps.append({"note": headset.hmd_serial, "tick": tick})
            headset.erase_records()
            instance.reset_history()
            return
    
    try:
        if not is_bootloader:
            notification.notify(f"{headset.hmd_serial} --> Connect!", f"(~˘▾˘)~ Updating... ::: FW v{headset.fw_version} --> v{target_firmware_version}", "Beyond Quick Updater", timeout=5, app_icon=icon_path_from_sn(headset.hmd_serial))
            log_info(headset.hmd_serial + ": " + f"Now updating HMD firmware to v{target_firmware_version}...")

            # stagger MCU updates
            headset.set_light_color(rgb=(255, 0, 255), led_mode=LEDModes.BLINKING)
            # update_timeslot += tick - start_tick  # offset granted timeslot by the time it took to get here
            # while update_timeslot > tick:
            #     time.sleep(1)

            headset.erase_records()
            instance.reset_history()
            headset.update_state = UpdateState.BOOTLOADER
            log_debug(f"{headset.hmd_serial}: BOOTLOADER ENTRY CMD for device at --> {target_device_path}")  # debug
            headset.enter_hmd_bootloader(confirm_bootloader=False)

        if is_bootloader:
            # stagger bootloader ops
            # update_timeslot += tick - start_tick  # offset granted timeslot by the time it took to get here
            # while update_timeslot > tick:
            #     time.sleep(1)

            log_info(headset.hmd_serial + ": " + f"Bootloader active. Proceeding with update...")
            log_debug(f"{headset.hmd_serial}: FW UPDATE WRITE for device at --> {target_device_path}")  # debug
            headset.update_hmd_firmware(firmware_payload)
            log_info(headset.hmd_serial + ": " + f"Update done. Reviving HMD @O@!!!")
            headset.erase_records(include_headsets_log=True)
            instance.reset_history()

            # notification.notify(f"{headset.hmd_serial} is updated", f"Complete!(づ ᴗ _ᴗ)づ♡ ::: FW is v{headset.fw_version} now~", "Beyond Quick Updater", timeout=10, app_icon=icon_path_from_sn(headset.hmd_serial))
            # log_info(headset.hmd_serial + f": HMD is now running firmware v{headset.fw_version}. Update complete!")
            # headset.update_state = UpdateState.DONE
            # success_count += 1
            # if CLEAR_HISTORY:
            #     instance.reset_history()
    except BaseException as e:
        print(e)
        log_warning(headset.hmd_serial + ": " + str(e))
        log_warning(headset.hmd_serial + ": Restarting root hub to try updating again...")
        headset.clear_hid_device()
        headset.erase_records()
        instance.reset_history()
        headset.restart_hub()


def reset_run():
    global success_count
    global fail_count
    global tick
    global update_timeslot_current
    global update_timestamps
    global matching

    if CLEAR_HISTORY:
        instance.reset_history()
        matching.reset()

    success_count = 0
    fail_count = 0
    update_timestamps = []
    # vvv likely not needed for quick updater vvv
    # tick = 0
    # lhcalib_timeslot_current = 0


def log_window():
    imgui.separator()
    hello_imgui.log_gui()
    time.sleep(0.1)



def log_loop():
    global window_title

    while True:
        immapp.run(log_window, window_title, with_markdown=True, window_restore_previous_geometry=True)


def tick_loop():
    global tick

    while True:
        time.sleep(1)
        with tick_lock:
            tick += 1


def clean_matched_paths_loop():
    global matching

    while True:
        time.sleep(15)
        matching.clean_matched_paths()


def device_cleanup_loop():
    global success_count, fail_count
    global last_clean_device_qty

    while True:
        # cleanup after n HMDs processed to keep the PC running smoothly
        while success_count + fail_count < last_clean_device_qty + USB_CLEANUP_THRESHOLD:
            time.sleep(5)

        log_info("Cleaning up USB devices...")
        bs_usb_tools.cleanup_beyond_devices(log_info, log_debug)
        last_clean_device_qty = success_count + fail_count


def device_queue_processor():
    """Worker thread that processes devices from the queue"""
    global matching
    while True:
        try:
            dev, i, new_devs, prev_devs = device_queue.get(timeout=1)
            try:
                matching.process_device(dev, i, new_devs, prev_devs)
            except Exception as e:
                log_warning(f"Error processing device: {e}")
            finally:
                device_queue.task_done()
        except queue.Empty:
            continue


def main():
    global tick
    global success_count
    global fail_count
    global update_timestamps
    global update_timestamps_lock
    global matching
    global target_firmware_path, target_firmware_version, firmware_payload
    global target_fpga_firmware_path, target_fpga_firmware_version, fpga_firmware_payload

    update_window_title()
    threading.Thread(target=tick_loop, daemon=True).start()
    threading.Thread(target=log_loop, daemon=True).start()

    if CLEAR_HISTORY:
        instance.reset_history()

    log_info(f"Beyond Quick Updater Tool ::: v{QUICK_UPDATER_VERSION}")
    log_info("IMPORTANT: Close the Bigscreen Beyond Utility (if it is running) to avoid issues!")

    target_firmware_path = get_target_firmware_path()
    target_firmware_version = bs_hmd_tools.get_file_fw_version(target_firmware_path)
    firmware_payload = bs_hmd_tools.load_beyondfw_file(target_firmware_path)
    log_info(f"Preloaded target firmware v{target_firmware_version} from {target_firmware_path}")

    target_fpga_firmware_path = get_target_fpga_firmware_path()
    target_fpga_firmware_version = bs_hmd_tools.get_dfu_file_fw_version(target_fpga_firmware_path)
    fpga_firmware_payload = bs_usb_tools.load_file(target_fpga_firmware_path)
    log_info(f"Preloaded target FPGA firmware v{target_fpga_firmware_version} from {target_fpga_firmware_path}")

    matching = HeadsetHIDMatching(lambda headset, hmd_update_timeslot: print(headset))
    instance.register_on_new_hid_device(matching.on_new_hid_device)
    
    # Start device queue processor thread
    threading.Thread(target=device_cleanup_loop, daemon=True).start()
    threading.Thread(target=device_queue_processor, daemon=True).start()
    threading.Thread(target=clean_matched_paths_loop, daemon=True).start()
    
    notification.notify(f"(っ´ཀ`)っ Quick Updater ::: v{QUICK_UPDATER_VERSION}", f"Connected HMDs will automatically be updated. IMPORTANT: Make sure the Beyond Utility is not running!", "Beyond Quick Updater", timeout=10, app_icon=icon_path_clear)
    log_info("----- Firmware updater ready! Connect the headsets you wish to update ------")

    futures = []
    with ThreadPoolExecutor(max_workers=MAX_THREADS) as tp:
        matching.on_match = lambda headset, hmd_update_timeslot: futures.append(tp.submit(update_hmd_firmware, headset, hmd_update_timeslot))

        while True:
            time.sleep(1)

    # reset_run()


if __name__ == '__main__':
    try:
        main()
    except BaseException as e:
        print(e)
    k = input("press close to exit")