import usb.core
import usb.util
import json
import os
import subprocess
import re

def enum_usb_devices():
    devices = usb.core.find(find_all=True)
    result = []
    for dev in devices:
        vendor_id = hex(dev.idVendor)
        product_id = hex(dev.idProduct)
        
        # 获取设备名称，部分设备可能没有字符串描述符
        try:
            device_name = usb.util.get_string(dev, dev.iProduct) or "Unknown"  # 必须以sudo的方式运行脚本
        except Exception:
            device_name = "Unknown"
        
        # 获取设备路径
        try:
            bus = dev.bus
            address = dev.address
            device_path = f"/dev/bus/usb/{bus:03d}/{address:03d}"
        except Exception:
            device_path = "Unknown"
        
        result.append({
            "deviceName": device_name,
            "vendorID": vendor_id,
            "productID": product_id,
            "devicePath": device_path,
            "subSystem": "USB"
        })
    return json.dumps(result, indent=2, ensure_ascii=False)


def enum_pcie_devices():
    result = []
    pci_devices_path = "/sys/bus/pci/devices"
    
    if not os.path.exists(pci_devices_path):
        return json.dumps(result, indent=2, ensure_ascii=False)
    
    for device_dir in os.listdir(pci_devices_path):
        device_path = os.path.join(pci_devices_path, device_dir)
        
        try:
            # 读取vendor ID和product ID
            with open(os.path.join(device_path, "vendor"), "r") as f:
                vendor_id = f.read().strip()
            
            with open(os.path.join(device_path, "device"), "r") as f:
                product_id = f.read().strip()
            
            # 获取设备名称
            device_name = "Unknown"
            try:
                # 尝试从modalias获取设备信息
                modalias_path = os.path.join(device_path, "modalias")
                if os.path.exists(modalias_path):
                    with open(modalias_path, "r") as f:
                        modalias = f.read().strip()
                
                # 尝试从driver获取驱动名称作为设备名称
                driver_path = os.path.join(device_path, "driver")
                if os.path.islink(driver_path):
                    driver_name = os.path.basename(os.readlink(driver_path))
                    device_name = f"PCIe Device ({driver_name})"
                
                # 尝试从subsystem_device和subsystem_vendor获取更详细信息
                subsystem_device_path = os.path.join(device_path, "subsystem_device")
                if os.path.exists(subsystem_device_path):
                    with open(subsystem_device_path, "r") as f:
                        subsystem_device = f.read().strip()
                    if device_name == "Unknown":
                        device_name = f"PCIe Device ({subsystem_device})"
                
            except Exception:
                device_name = "Unknown"
            
            # 构建设备路径 (PCI设备路径格式)
            pci_device_path = f"/sys/bus/pci/devices/{device_dir}"
            
            result.append({
                "deviceName": device_name,
                "vendorID": vendor_id,
                "productID": product_id,
                "devicePath": pci_device_path,
                "subSystem": "PCIe"
            })
            
        except Exception as e:
            # 如果读取某个设备信息失败，跳过该设备
            continue
    
    return json.dumps(result, indent=2, ensure_ascii=False)


def checkPid(devPath, logger):
    processPid = -1
    # 通过device_path查看当前占用该设备的进程pid号  
    try:
        # 首先检查设备路径是否存在
        if not os.path.exists(devPath):
            return processPid
        
        # 添加-w参数禁用警告，使用2>/dev/null重定向标准错误
        result = subprocess.run(['lsof', '-w', devPath], 
                               stdout=subprocess.PIPE, 
                               stderr=subprocess.DEVNULL, 
                               text=True)
        
        if result.returncode == 0:
            output = result.stdout.strip()
            if output:
                # 使用正则表达式提取PID
                # lsof输出格式: COMMAND PID USER FD TYPE DEVICE SIZE/OFF NODE NAME
                pid_match = re.search(r'^\S+\s+(\d+)\s+', output, re.MULTILINE)
                if pid_match:
                    processPid = int(pid_match.group(1))
                    logger.debug(f"Device {devPath} is currently used by process with PID: {processPid}")
            
    except FileNotFoundError:
        logger.error("lsof command not found. Please install lsof: sudo apt-get install lsof")
    except Exception as e:
        logger.error(f"An error occurred while checking the device: {str(e)}")
    
    return processPid