import os

gpu_class = 0x030000
pci_path = "/sys/bus/pci/devices"
misc_path = "/usr/share/misc/pci.ids"


class GPU:
    def __init__(self, vendor_id, vendor, device_id, device, driver):
        self.vendor = vendor
        self.device = device
        self.vendor_id = vendor_id
        self.device_id = device_id
        self.driver = driver


def int2hex(num):
    return str(hex(num)[2:]).upper()


def parse_pci_ids(file_path):
    with open(file_path, "r") as f:
        lines = f.readlines()

    devices = {}
    current_vendor = None

    for line in lines:
        if line.startswith("#") or line.strip() == "":
            continue

        if not line.startswith("\t"):
            # This is a vendor entry
            if current_vendor:
                devices[current_vendor] = vendor_info
            vendor_id, vendor_name = line.strip().split(" ", 1)
            current_vendor = int2hex(int(vendor_id, 16))
            vendor_info = {"vendor_name": vendor_name.strip(), "devices": []}
        else:
            # This is a device entry
            device_id, device_name = line.strip().split(" ", 1)
            vendor_info["devices"].append(
                {
                    "device_id": int2hex(int(device_id, 16)),
                    "device_name": device_name.strip(),
                    "vendor_name": vendor_info["vendor_name"],
                }
            )
    if current_vendor:
        devices[current_vendor] = vendor_info
    return devices


parsed_pci_ids = parse_pci_ids(misc_path)


def get():
    gpus = []
    for root, dirs, files in os.walk(pci_path):
        for pci_dir in dirs:
            dev_content_path = os.path.join(root, pci_dir, "class")
            if os.path.exists(dev_content_path):
                with open(dev_content_path, "r") as f:
                    cont = f.readline().strip()
                    cont = int(cont, 16)
                    if cont == gpu_class:
                        gpu_dev = None
                        gpu_ven = None
                        gpu_drv = None
                        gpu_dev = cont
                        ven_content_path = os.path.join(root, pci_dir, "vendor")
                        if os.path.exists(ven_content_path):
                            with open(ven_content_path) as f:
                                gpu_ven = f.readline().strip()
                                gpu_ven = int2hex(int(gpu_ven, 16))

                        dev_content_path = os.path.join(root, pci_dir, "device")
                        if os.path.exists(dev_content_path):
                            with open(dev_content_path) as f:
                                gpu_dev = f.readline().strip()
                                gpu_dev = int2hex(int(gpu_dev, 16))

                        drv_content_path = os.path.join(
                            root, pci_dir, "driver", "module"
                        )
                        if os.path.exists(drv_content_path):
                            gpu_drv_p = os.readlink(drv_content_path)
                            gpu_drv = os.path.basename(gpu_drv_p)

                        gpu_ven_name = parsed_pci_ids[gpu_ven]["vendor_name"]
                        gpu_dev_name = None
                        for dev in parsed_pci_ids[gpu_ven]["devices"]:
                            if dev["device_id"] == gpu_dev:
                                gpu_dev_name = dev["device_name"]

                        gpus.append(
                            GPU(gpu_ven, gpu_ven_name, gpu_dev, gpu_dev_name, gpu_drv)
                        )
    return gpus
