#!/usr/bin/env python3
"""
Checks for Nvidia GPUs, Broadcom network cards, and AMD GPUs,
and returns a dictionary with the hardware description and required packages,
including support for legacy hardware.

Legacy hardware support is provided by parsing the Linux Kernel Device Database (LKDDb)

Author: Arjen Balfoort 
Date: 2026-01-15
Dependencies: nvidia-detect, (pciutils)
"""

import subprocess
import re
import urllib.request
import os
import glob
import shutil
from typing import Optional
from utils import DebSoftwareApp


class LKDDbParser:
    """
    Parses the Linux Kernel Device Database (LKDDb) list file.
    Downloads the lkddb.list file if the kernel version has changed since the last download.
    Extracts device IDs for specific vendors and drivers.

    Args:
        save_directory (str): Directory to save the lkddb.list file, kernel.txt, and log files.
    
    Attributes:
        url (str): URL to download the lkddb.list file.
        filename (str): Local filename for the lkddb.list file.
        kernel_version_file (str): Filename to store the current kernel version.
        current_kernel_version (str): The current kernel version.
        saved_kernel_version (str): The saved kernel version from the last download.
        amd_legacy_radeon (list): List of AMD legacy Radeon device IDs.
        amd_legacy_r128 (list): List of AMD legacy R128 device IDs.
        amd_legacy_mach64 (list): List of AMD legacy Mach64 device IDs.
        broadcom_legacy (list): List of Broadcom legacy device IDs.
        broadcom_legacy_brcm80211 (list): List of Broadcom legacy brcm80211 device IDs.
    """
    def __init__(self, app: DebSoftwareApp):
        self.app: DebSoftwareApp = app
        self.lkddb_list: str = os.path.join(self.app.save_dir, "lkddb.list")
        self.kernel_txt: str = os.path.join(self.app.save_dir, "kernel.txt")

        # Get current and saved kernel versions
        self.current_kernel_version: str = self.get_current_kernel_version()
        self.saved_kernel_version: Optional[str] = self.read_saved_kernel_version()

        # Clean up old log files
        for log in glob.glob(os.path.join(self.app.save_dir, 'lkddb.list-*.log')):
            os.remove(log)

        # Download or use existing lkddb.list
        self.get_lkddb()

        # Extract device IDs for specific vendors and drivers
        self.amd_legacy_radeon: list[str] = self._get_device_ids(
            hw_vendor_id=self.app.amd_vendor_id,
            driver="radeon_drv"
        )
        self.amd_legacy_r128: list[str] = self._get_device_ids(
            hw_vendor_id=self.app.amd_vendor_id,
            driver="aty128fb"
        )
        self.amd_legacy_mach64: list[str] = self._get_device_ids(
            hw_vendor_id=self.app.amd_vendor_id,
            driver="atyfb_base"
        )
        self.broadcom_legacy: list[str] = self._get_device_ids(
            hw_vendor_id=self.app.broadcom_vendor_id,
            driver="b43_pci_bridge"
        )
        self.broadcom_legacy_brcm80211: list[str] = self._get_device_ids(
            hw_vendor_id=self.app.broadcom_vendor_id,
            driver="brcm80211"
        )

    @staticmethod
    def get_current_kernel_version():
        """ Get the current kernel version using uname -r. """
        return subprocess.check_output(["uname", "-r"], universal_newlines=True).strip()

    def _save_kernel_version(self):
        """ Save the current kernel version to kernel.txt. """
        with open(file=self.kernel_txt, mode="w", encoding="utf-8") as f:
            f.write(self.current_kernel_version)
        self.saved_kernel_version = self.current_kernel_version

    def read_saved_kernel_version(self):
        """ Read the saved kernel version from kernel.txt. """
        if os.path.exists(self.kernel_txt):
            with open(file=self.kernel_txt, mode="r", encoding="utf-8") as f:
                return f.read().strip()
        return None

    def get_lkddb(self, force_download: bool = False):
        """
        Download lkddb.list only if kernel version changed or file does not exist.
        Args:
            force_download (bool): If True, forces re-download of the file."""
        if force_download:
            if os.path.exists(self.lkddb_list):
                os.remove(self.lkddb_list)
            if os.path.exists(self.kernel_txt):
                os.remove(self.kernel_txt)

        if self.saved_kernel_version != self.current_kernel_version:
            self.app.logging.info("Downloading %s...", self.app.lkddb_url)
            try:
                urllib.request.urlretrieve(self.app.lkddb_url, self.lkddb_list)
                self._save_kernel_version()
            except (urllib.error.URLError,
                    urllib.error.HTTPError,
                    urllib.error.ContentTooShortError) as e:
                self.app.logging.error("Unexpected error: %s", e)
                if os.path.exists(self.lkddb_list):
                    self.app.logging.warning("Using local %s instead.", self.lkddb_list)
                else:
                    # Save drivers data if lkddb.list does not exist in the save directory
                    lkdbb = os.path.join(self.app.data_dir, 'lkddb.list')
                    if os.path.exists(lkdbb):
                        shutil.copyfile(lkdbb, self.lkddb_list)

                    # Check if lkddb.list exists
                    if not os.path.exists(self.lkddb_list):
                        raise RuntimeError(f"Failed to download or use local {self.lkddb_list}") from e
        else:
            self.app.logging.info("Kernel version unchanged. Using local %s.", self.lkddb_list)

    def _get_device_ids(self, hw_vendor_id: str, driver: str) -> list[str]:
        """
        Extract device IDs for a given vendor and driver using regex.
        Args:
            hw_vendor_id (str): The hardware vendor ID to search for.
            driver (str): The driver name to search for.
        Returns:
            list: A list of device IDs as strings.
        """
        if os.path.exists(self.lkddb_list):
            self.app.logging.info("Parsing %s for vendor %s and driver %s...",
                         self.lkddb_list, hw_vendor_id, driver)
            device_ids = []
            pattern = re.compile(rf'\s+{hw_vendor_id}\s+([0-9a-z]{{4}}).*{driver}.*')
            with open(file=self.lkddb_list, mode="r", encoding="utf-8", errors="ignore") as f:
                for line in f:
                    match = pattern.search(line)
                    if match:
                        device_id = match.group(1)
                        if device_id and device_id not in device_ids:
                            device_ids.append(device_id)

            # Write device IDs to a file for reference
            with open(file=f"{self.lkddb_list}-{hw_vendor_id}.log",
                      mode='a',
                      encoding="utf-8") as f:
                f.write('\n'.join(str(id) for id in device_ids) + '\n\n')

            return device_ids
        return []


class ProprietaryHardwareChecker():
    """
    A threaded class for detecting proprietary hardware on Debian-based systems.
    This class checks for Nvidia GPUs, Broadcom network cards, and AMD GPUs,
    and returns a dictionary with the hardware description and required packages,
    including support for legacy hardware.

    Attributes:
        result (dict): Stores the detection results.
        amd_legacy (dict): Maps AMD device IDs to driver packages (legacy).
        broadcom_legacy (dict): Maps Broadcom device IDs to driver packages (legacy).
    """

    def __init__(self,
                 app: DebSoftwareApp,
                 test: bool=False) -> None:
        """
        Initializes the ProprietaryHardwareChecker.
        """
        self.app = app
        self.result = {}

        # AMD device IDs and their driver packages (modern and legacy)
        self.amd_drivers = ['firmware-amd-graphics']
        self.amd_legacy_radeon_drivers = ['xserver-xorg-video-ati']
        self.amd_legacy_r128_drivers = ['xserver-xorg-video-r128']
        self.amd_legacy_mach64_drivers = ['xserver-xorg-video-mach64']

        # Broadcom device IDs and their driver packages (modern and legacy)
        self.broadcom_drivers = ['firmware-b43-installer', 'broadcom-sta-dkms']
        self.broadcom_legacy_drivers = ['firmware-b43legacy-installer']
        self.broadcom_legacy_brcm80211_drivers = ['firmware-brcm80211']

        # For testing purposes, predefined lspci output
        self.test = test
        self._test_hardware = [
            '06:00.0 VGA compatible controller [0300]: Advanced Micro Devices, Inc. '
                '[AMD/ATI] Rembrandt [Radeon 680M] [1002:1681] (rev c7)',
            '06:00.0 VGA compatible controller [0300]: NVIDIA Corporation GA106 '
                '[GeForce RTX 3060 Lite Hash Rate] [10de:2504] (rev a1)',
            '01:00.0 Network controller [0280]: Broadcom Inc. and subsidiaries '
                'BCM43142 802.11b/g/n [14e4:4365] (rev 01)'
        ]

        # Initialize LKDD parser to get legacy device IDs
        try:
            self.lkddb_parser = LKDDbParser(self.app)
        except Exception as e:
            self.app.logging.error("Error initializing LKDDbParser: %s", e)
            return

        # Start hardware detection
        self.detect_hardware()

    def detect_hardware(self) -> None:
        """
        Detects proprietary hardware using lspci and determines required packages.
        """
        if self.test:
            self.app.logging.info("Using test data for hardware detection.")
            output = self._test_hardware
        else:
            # Run lspci to list all PCI devices with vendor and device IDs
            try:
                self.app.logging.info("Running lspci to list hardware...")
                output = subprocess.check_output(['lspci', '-nn'],
                                                stderr=subprocess.PIPE).decode('utf-8').splitlines()
            except subprocess.CalledProcessError as e:
                self.app.logging.error("Error running lspci: %s", e)
                return

        for line in output:
            match = re.search(
                rf'(.*)\]:\s+(.*)\s+\[({self.app.amd_vendor_id}|{self.app.broadcom_vendor_id}|'
                rf'{self.app.nvidia_vendor_id}):([a-z0-9]{{4}})\]', line
            )
            if match:
                # Extract description, vendor and device IDs from the line
                hd_info = match.group(1).lower()
                description = match.group(2).strip()
                hw_vendor_id = match.group(3)
                hw_device_id = match.group(4)

                # Check for supported AMD GPUs
                if hw_vendor_id == self.app.amd_vendor_id and 'vga' in hd_info:
                    self.app.logging.info("AMD GPU detected: %s", description)
                    packages = self.amd_drivers
                    # Check for legacy Radeon drivers
                    if hw_device_id in self.lkddb_parser.amd_legacy_radeon:
                        packages = self.amd_legacy_radeon_drivers
                        self.app.logging.info("Legacy Radeon drivers needed")
                    elif hw_device_id in self.lkddb_parser.amd_legacy_r128:
                        packages = self.amd_legacy_r128_drivers
                        self.app.logging.info("Legacy R128 drivers needed")
                    elif hw_device_id in self.lkddb_parser.amd_legacy_mach64:
                        packages = self.amd_legacy_mach64_drivers
                        self.app.logging.info("Legacy Mach64 drivers needed")
                    packages += self._get_additional_packages(hw_vendor_id)

                    # Filter packages on not installed status
                    self.app.logging.info("AMD packages needed: %s", ', '.join(packages))
                    packages = self._filter_packages(packages)
                    self.app.logging.info(
                        "AMD packages to install after check: %s", ', '.join(packages)
                    )

                    # Check if the proprietary driver is installed by AMD installer
                    # Empty packages list if driver is installed
                    if len(packages) > 0:
                        if set(self.amd_drivers).intersection(packages) and \
                        self._is_proprietary_driver_installed(vendor_id=self.app.amd_vendor_id):
                            packages = []
                            self.app.logging.info("AMD proprietary driver is installed")

                    self.result[self.app.amd_vendor_id] = {
                        'device_id': hw_device_id,
                        'name': description,
                        'packages': packages
                    }

                # Check for supported Broadcom network cards
                elif hw_vendor_id == self.app.broadcom_vendor_id:
                    self.app.logging.info("Broadcom network card detected: %s", description)
                    packages = self.broadcom_drivers
                    # Check for legacy Broadcom drivers
                    if hw_device_id in self.lkddb_parser.broadcom_legacy:
                        packages = self.broadcom_legacy_drivers
                        self.app.logging.debug("Legacy Broadcom drivers needed")
                    elif hw_device_id in self.lkddb_parser.broadcom_legacy_brcm80211:
                        packages = self.broadcom_legacy_brcm80211_drivers
                        self.app.logging.debug("Legacy Broadcom brcm80211 drivers needed")
                    packages += self._get_additional_packages(hw_vendor_id)

                    # Filter packages on not installed status
                    self.app.logging.info(
                        "Broadcom packages needed: %s", ', '.join(packages)
                    )
                    packages = self._filter_packages(packages)
                    self.app.logging.info(
                        "Broadcom packages to install after check: %s", ', '.join(packages)
                    )

                    # Check if the proprietary driver is installed by Broadcom installer
                    # Empty packages list if driver is installed
                    if len(packages) > 0:
                        if set(self.broadcom_drivers).intersection(packages) and \
                        self._is_proprietary_driver_installed(
                            vendor_id=self.app.broadcom_vendor_id
                        ):
                            packages = []
                            self.app.logging.info(
                                "Broadcom proprietary driver is installed"
                            )

                    self.result[self.app.broadcom_vendor_id] = {
                        'device_id': hw_device_id,
                        'name': description,
                        'packages': packages
                    }
                # Check for Nvidia GPUs
                elif hw_vendor_id == self.app.nvidia_vendor_id and 'vga' in hd_info:
                    self.app.logging.info("Nvidia GPU detected: %s", description)
                    driver = self._check_nvidia_packages()
                    if driver is not None:
                        packages = self._filter_packages(
                            [driver, 'firmware-misc-nonfree'] + \
                            self._get_additional_packages(hw_vendor_id)
                        )

                        self.app.logging.info("Nvidia packages needed: %s", ', '.join(packages))

                        # Check if the proprietary driver is installed by Nvidia installer
                        # Empty packages list if driver is installed
                        if len(packages) > 0:
                            if self._is_proprietary_driver_installed(
                                vendor_id=self.app.nvidia_vendor_id
                            ):
                                packages = []
                                self.app.logging.info("Nvidia proprietary driver is installed")

                        self.result[self.app.nvidia_vendor_id] = {
                            'device_id': hw_device_id,
                            'name': description,
                            'packages': packages
                        }

    def _get_additional_packages(self, hw_vendor_id: str) -> list[str]:
        """
        Get additional packages required for building drivers. 
        
        Args:
            hw_vendor_id (str): The vendor ID of the hardware.
            Returns:
            list: A list of additional package names.
        """
        additional_packages = [f"linux-headers-{self.lkddb_parser.current_kernel_version}",
                               "firmware-linux-nonfree"]

        if hw_vendor_id in [self.app.nvidia_vendor_id, self.app.amd_vendor_id]:
            return additional_packages + ["build-essential"]
        return additional_packages + ["wireless-tools"]

    def _check_nvidia_packages(self) -> Optional[str]:
        """
        Uses nvidia-detect to determine the correct driver package for Nvidia GPUs.
        """
        self.app.logging.info("Run nvidia-detect to get driver package")
        try:
            # Run nvidia-detect and capture output
            nvidia_detect = subprocess.run(
                ['nvidia-detect'],
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                check=True,
                text=True
            ).stdout.splitlines()

            # Filter lines containing 'nvidia-' and remove spaces
            package = None
            for line in nvidia_detect:
                self.app.logging.debug("nvidia-detect line: %s", line)
                if 'nvidia-' in line:
                    package = line.strip()
                    self.app.logging.info("Nvidia driver package to install: %s", package)
                    break

            return package

        except subprocess.CalledProcessError as e:
            self.app.logging.error("Error running nvidia-detect: %s", e)
            return None

    def _filter_packages(self, packages: list[str]) -> list[str]:
        """
        Filters packages based on not installed status.

        Args:
            packages (list): A list of packages to filter.

        Returns:
            list: A list of packages not yet installed.
        """
        filtered_packages = []

        for package in packages:
            try:
                pck_result = subprocess.run(
                    ["dpkg", "-l", package],
                    capture_output=True,
                    text=True,
                    check=False
                ).stdout.strip()
                if "ii" not in pck_result:
                    filtered_packages.append(package)
                else:
                    self.app.logging.debug("Package %s is already installed", package)
            except subprocess.CalledProcessError:
                continue

        return filtered_packages

    def _is_proprietary_driver_installed(self, vendor_id: str) -> bool:
        """
        Check if a drm driver is installed by any other means then debian packages.

        Args:
            vendor_id (str): The vendor ID of the hardware.
        
        Returns:
            bool: True if the driver is installed, False otherwise.
        """
        if vendor_id == self.app.nvidia_vendor_id:
            pattern = 'nvidia-drm.ko*'
        elif vendor_id == self.app.broadcom_vendor_id:
            pattern = 'wl.ko*'
        elif vendor_id == self.app.amd_vendor_id:
            try:
                result = subprocess.run(
                    ["dpkg", "-l"],
                    capture_output=True,
                    text=True,
                    check=True
                )
                return "amdgpu-dkms" in result.stdout
            except subprocess.CalledProcessError:
                return False
        else:
            return False

        # It is not installed: check if owned by Nvidia (.run installer)
        drivers = glob.glob(f"/lib/modules/{self.lkddb_parser.current_kernel_version}"
                            f"/**/{pattern}", recursive=True)
        if len(drivers) > 0:
            self.app.logging.debug("Driver %s is installed", pattern)
            return True
        self.app.logging.debug("Driver %s is not installed", pattern)
        return False
