import os, sys, subprocess

# Get the base path for bundled files
if getattr(sys, 'frozen', False):
    _BASE_PATH = sys._MEIPASS
else:
    _BASE_PATH = os.path.dirname(os.path.abspath(__file__))

DCC_PATH = os.path.join(_BASE_PATH, 'tools', 'DeviceCleanupCmd.exe')
LUD_PATH = os.path.join(_BASE_PATH, 'tools', 'ListUsbDevs.exe')
RUP_PATH = os.path.join(_BASE_PATH, 'tools', 'RestartUsbPort.exe')
DFU_PATH = os.path.join(_BASE_PATH, 'tools', 'dfu-util.exe')
DELAY_AFTER_RESTART = 5  # seconds to wait after restarting a USB port


def get_info(search_str, as_list=False):
  key_count = {}
  device_list = []

  assert os.path.exists(LUD_PATH), "Path to ListUsbDevs.exe does not exist"
  cmd = [LUD_PATH, '-na', search_str]
  print("Running command: " + ' '.join(cmd))
  p = subprocess.Popen(cmd, stdout=subprocess.PIPE, stdin=subprocess.PIPE, stderr=subprocess.PIPE)

  rslt = None
  while rslt is None:
      rslt = p.poll()
      line = p.stdout.readline().decode().rstrip()
      if line == "":
          continue
      print(line)
      if "No matching " in line:
          print("\r\nNo matching devices found")
          return None
      if ": " in line: # parse key/value pairs, removing whitespace from the ends
          key, value = line.split(": ", 1)
          key_count[key] = key_count.get(key, 0) + 1
          device_index = key_count[key] - 1

          if device_index >= len(device_list):
              device_list.append({})
          if "USB Port" in key:
              device_list[device_index]["PortChain"] = value.strip().split("-")
          else:
              device_list[device_index][key.strip()] = value.strip()

  assert rslt == 0, p.stderr.readline().decode().rstrip()
  if as_list:
      return device_list
  return device_list[0] if device_list else None


def restart_port(search_str):
  assert os.path.exists(RUP_PATH), "Path to RestartUsbPort.exe does not exist"
  cmd = [RUP_PATH, '-n', '-na', f"-w:{DELAY_AFTER_RESTART * 1000}", search_str]
  print("Running command: " + ' '.join(cmd))
  p = subprocess.Popen(cmd, stdout=subprocess.PIPE, stdin=subprocess.PIPE, stderr=subprocess.PIPE)

  ret = False
  rslt = None
  while rslt is None:
      rslt = p.poll()
      line = p.stdout.readline().decode().rstrip()
      if line == "":
          continue
      print(line)
      if "success" in line: # parse key/value pairs, removing whitespace from the ends
          ret = True

  assert rslt == 0, p.stderr.readline().decode().rstrip()
  return ret


def cleanup_beyond_devices(log_info, log_debug, exclude_individual_reports=True):
    log_info("Performing device cleanup")
    cmd = [DCC_PATH, '-na', '-s', "SWD\\MMDEVAPI*0.0.1*", "SWD\\MMDEVAPI*0.0.0*", "*VID_28DE*PID_2102*", "*VID_28DE*PID_2300*", "*VID_0424*PID_3803*", "*VID_35BD*"]
    print("Opening command: " + ' '.join(cmd))
    p = subprocess.Popen(cmd, stdout=subprocess.PIPE, stdin=subprocess.PIPE, stderr=subprocess.PIPE)

    rslt = None
    while rslt is None:
        rslt = p.poll()
        line = p.stdout.readline().decode().rstrip()
        if line == "" or "OK" in line:
            continue
        print(line)
        if "removing device" not in line or not exclude_individual_reports:
            log_debug("DEVICE CLEANUP: " + line)
        if "Removed " in line:
            print("\r\n")
            log_info("DEVICE CLEANUP: " + line)
    assert rslt == 0, p.stderr.readline().decode().rstrip()


def dfu_format_port_chain(port_chain):
    if isinstance(port_chain, str):
        port_chain = port_chain.replace('.', '-').split('-')
    # Reconstruct port chain string with '-' as the first separator, '.' for subsequent
    return '-'.join(port_chain[0:2]) + '.'.join([''] + port_chain[2:])


def dfu_device_exists(dfu_vid, dfu_pid, dfu_port_chain=None):
    """
    Confirm whether dfu-util.exe can detect the target DFU device
    
    Args:
        dfu_vid: VID of the target DFU device (e.g., "28DE")
        dfu_pid: PID of the target DFU device (e.g., "2102")
        dfu_port_chain: Optional port chain location (e.g., ['1', '2', '3', '4'] or "1-2-3-4")
        
    Returns:
        bool: True if device is detected, False otherwise
    """
    assert os.path.exists(DFU_PATH), "Path to dfu-util.exe does not exist"
    
    # Build command
    cmd = [DFU_PATH, '-d', f",{dfu_vid}:{dfu_pid}", '-l']
    
    # Add port chain if provided
    if dfu_port_chain is not None:
        dfu_port_chain = dfu_format_port_chain(dfu_port_chain)
        cmd.insert(3, '-p')
        cmd.insert(4, dfu_port_chain)
    
    print("Running command: " + ' '.join(cmd))
    
    # Run dfu-util with firmware payload sent to stdin
    p = subprocess.Popen(
        cmd,
        stdout=subprocess.PIPE,
        stdin=subprocess.PIPE,
        stderr=subprocess.PIPE
    )
    
    # Send firmware payload to stdin
    try:
        stdout, stderr = p.communicate()
    except Exception as e:
        print(f"Error communicating with dfu-util: {e}")
        return False
    
    # Print output
    if stdout:
        line_ignore_list = [
            "dfu-util ",
            "Copyright",
            "This program is ",
            "Please report bugs ",
        ]
        for line in stdout.decode().splitlines():
            if line == "":
                continue
            if line and not any(ign in line for ign in line_ignore_list):
                print(line)
                if "Found DFU" in line:
                    return True
    
    if stderr:
        for line in stderr.decode().splitlines():
            print(line)
    
    return False


def load_file(file_path: str) -> bytearray:
    """Load a file. Intended for use with .dfu files."""
    with open(file_path, 'rb') as f:
        data = f.read()

    return bytearray(data)


def dfu_download(firmware_payload, dfu_vid, dfu_pid, dfu_port_chain=None):
    """
    Download firmware to a DFU device using dfu-util.exe
    
    Args:
        firmware_payload: bytearray containing the entire .dfu file. Create using load_file()
        dfu_vid: VID of the target DFU device (e.g., "28DE")
        dfu_pid: PID of the target DFU device (e.g., "2102")
        dfu_port_chain: Optional port chain location (e.g., ['1', '2', '3', '4'] or "1-2-3-4")
        
    Returns:
        bool: True if download succeeded, False otherwise
    """
    assert os.path.exists(DFU_PATH), "Path to dfu-util.exe does not exist"
    
    # Build command
    cmd = [DFU_PATH, '-d', f",{dfu_vid}:{dfu_pid}", '-D', '-']
    
    # Add port chain if provided
    if dfu_port_chain is not None:
        dfu_port_chain = dfu_format_port_chain(dfu_port_chain)
        cmd.insert(3, '-p')
        cmd.insert(4, dfu_port_chain)
    
    print("Running command: " + ' '.join(cmd))
    
    # Run dfu-util with firmware payload sent to stdin
    p = subprocess.Popen(
        cmd,
        stdout=subprocess.PIPE,
        stdin=subprocess.PIPE,
        stderr=subprocess.PIPE
    )
    
    # Send firmware payload to stdin
    try:
        stdout, stderr = p.communicate(input=bytes(firmware_payload))
    except Exception as e:
        print(f"Error communicating with dfu-util: {e}")
        return False
    
    # Print output
    if stdout:
        line_ignore_list = [
            "dfu-util ",
            "Copyright",
            "This program is ",
            "Please report bugs ",
            "Download\t",
        ]
        for line in stdout.decode().splitlines():
            if line == "":
                continue
            if line and not any(ign in line for ign in line_ignore_list):
                print(line)
    
    if stderr:
        for line in stderr.decode().splitlines():
            print(line)
    
    # Check return code
    if p.returncode == 0:
        print("DFU download completed successfully")
        return True
    else:
        print(f"DFU download failed with return code: {p.returncode}")
        return False