# Hardware Detection Logic for Root Patching
# Returns a dictionary of patches with boolean values
# Used when supplying data to sys_patch.py
# Copyright (C) 2020-2022, Dhinak G, Mykola Grymalyuk

from resources import constants, device_probe, utilities
from data import model_array, os_data, sip_data, sys_patch_dict

class detect_root_patch:
    def __init__(self, model, versions):
        self.model = model
        self.constants: constants.Constants() = versions
        self.computer = self.constants.computer

        # GPU Patch Detection
        self.nvidia_legacy= False
        self.kepler_gpu= False
        self.amd_ts1= False
        self.amd_ts2= False
        self.iron_gpu= False
        self.sandy_gpu= False
        self.ivy_gpu= False

        # Misc Patch Detection
        self.brightness_legacy= False
        self.legacy_audio= False
        self.legacy_wifi= False
        self.legacy_gmux= False
        self.legacy_keyboard_backlight= False

        # Patch Requirements
        self.amfi_must_disable= False
        self.supports_metal= False

        # Validation Checks
        self.sip_enabled = False
        self.sbm_enabled = False
        self.amfi_enabled = False
        self.fv_enabled = False
        self.dosdude_patched = False

    
    def detect_gpus(self):
        gpus = self.constants.computer.gpus
        non_metal_os = os_data.os_data.catalina
        for i, gpu in enumerate(gpus):
            if gpu.class_code and gpu.class_code != 0xFFFFFFFF:
                print(f"- Found GPU ({i}): {utilities.friendly_hex(gpu.vendor_id)}:{utilities.friendly_hex(gpu.device_id)}")
                if gpu.arch in [device_probe.NVIDIA.Archs.Tesla, device_probe.NVIDIA.Archs.Fermi]:
                    if self.constants.detected_os > non_metal_os:
                        self.nvidia_legacy = True
                        self.amfi_must_disable = True
                        self.legacy_keyboard_backlight = self.check_legacy_keyboard_backlight()
                elif gpu.arch == device_probe.NVIDIA.Archs.Kepler:
                    if self.constants.detected_os > os_data.os_data.big_sur:
                        # Kepler drivers were dropped with Beta 7
                        # 12.0 Beta 5: 21.0.0 - 21A5304g
                        # 12.0 Beta 6: 21.1.0 - 21A5506j
                        # 12.0 Beta 7: 21.1.0 - 21A5522h
                        if self.constants.detected_os == os_data.os_data.monterey and self.constants.detected_os_minor > 0:
                            if "21A5506j" not in self.constants.detected_os_build:
                                self.kepler_gpu = True
                                self.supports_metal = True
                elif gpu.arch == device_probe.AMD.Archs.TeraScale_1:
                    if self.constants.detected_os > non_metal_os:
                        self.amd_ts1 = True
                        self.amfi_must_disable = True
                elif gpu.arch == device_probe.AMD.Archs.TeraScale_2:
                    if self.constants.detected_os > non_metal_os:
                        self.amd_ts2 = True
                        self.amfi_must_disable = True
                elif gpu.arch == device_probe.Intel.Archs.Iron_Lake:
                    if self.constants.detected_os > non_metal_os:
                        self.iron_gpu = True
                        self.amfi_must_disable = True
                        self.legacy_keyboard_backlight = self.check_legacy_keyboard_backlight()
                elif gpu.arch == device_probe.Intel.Archs.Sandy_Bridge:
                    if self.constants.detected_os > non_metal_os:
                        self.sandy_gpu = True
                        self.amfi_must_disable = True
                        self.legacy_keyboard_backlight = self.check_legacy_keyboard_backlight()
                elif gpu.arch == device_probe.Intel.Archs.Ivy_Bridge:
                    if self.constants.detected_os > os_data.os_data.big_sur:
                        self.ivy_gpu = True
                        self.supports_metal = True
        if self.supports_metal is True:
            # Avoid patching Metal and non-Metal GPUs if both present, prioritize Metal GPU
            # Main concerns are for iMac12,x with Sandy iGPU and Kepler dGPU
            self.nvidia_legacy = False
            self.amd_ts1 = False
            self.amd_ts2 = False
            self.iron_gpu = False
            self.sandy_gpu = False
            self.legacy_keyboard_backlight = False
    
    def check_dgpu_status(self):
        dgpu = self.constants.computer.dgpu
        if dgpu:
            if dgpu.class_code and dgpu.class_code == 0xFFFFFFFF:
                # If dGPU is disabled via class-codes, assume demuxed
                return False
            return True
        return False

    def detect_demux(self):
        # If GFX0 is missing, assume machine was demuxed
        # -wegnoegpu would also trigger this, so ensure arg is not present
        if not "-wegnoegpu" in (utilities.get_nvram("boot-args") or ""):
            igpu = self.constants.computer.igpu
            dgpu = self.check_dgpu_status()
            if igpu and not dgpu:
                return True
        return False
    
    def check_legacy_keyboard_backlight(self):
        # iMac12,x+ have an 'ACPI0008' device, but it's not a keyboard backlight
        # Best to assume laptops will have a keyboard backlight
        if self.model.startswith("MacBook"):
            return self.constants.computer.ambient_light_sensor
        return False

    def detect_patch_set(self):
        self.detect_gpus()
        if self.model in model_array.LegacyBrightness:
            if self.constants.detected_os > os_data.os_data.catalina:
                self.brightness_legacy = True

        if self.model in ["iMac7,1", "iMac8,1"] or (self.model in model_array.LegacyAudio and utilities.check_kext_loaded("AppleALC", self.constants.detected_os) is False):
            # Special hack for systems with botched GOPs
            # TL;DR: No Boot Screen breaks Lilu, therefore breaking audio
            if self.constants.detected_os > os_data.os_data.catalina:
                self.legacy_audio = True

        if (
            isinstance(self.constants.computer.wifi, device_probe.Broadcom)
            and self.constants.computer.wifi.chipset in [device_probe.Broadcom.Chipsets.AirPortBrcm4331, device_probe.Broadcom.Chipsets.AirPortBrcm43224]
        ) or (isinstance(self.constants.computer.wifi, device_probe.Atheros) and self.constants.computer.wifi.chipset == device_probe.Atheros.Chipsets.AirPortAtheros40):
            if self.constants.detected_os > os_data.os_data.big_sur:
                self.legacy_wifi = True

        # if self.model in ["MacBookPro5,1", "MacBookPro5,2", "MacBookPro5,3", "MacBookPro8,2", "MacBookPro8,3"]:
        if self.model in ["MacBookPro8,2", "MacBookPro8,3"]:
            # Sierra uses a legacy GMUX control method needed for dGPU switching on MacBookPro5,x
            # Same method is also used for demuxed machines
            # Note that MacBookPro5,x machines are extremely unstable with this patch set, so disabled until investigated further
            # Ref: https://github.com/dortania/OpenCore-Legacy-Patcher/files/7360909/KP-b10-030.txt
            if self.constants.detected_os > os_data.os_data.high_sierra:
                if self.model in ["MacBookPro8,2", "MacBookPro8,3"]:
                    # Ref: https://doslabelectronics.com/Demux.html
                    if self.detect_demux() is True:
                        self.legacy_gmux = True
                else:
                    self.legacy_gmux = True
        
        self.root_patch_dict = {
            "Graphics: Nvidia Tesla": self.nvidia_legacy,
            "Graphics: Nvidia Kepler": self.kepler_gpu,
            "Graphics: AMD TeraScale 1": self.amd_ts1,
            "Graphics: AMD TeraScale 2": self.amd_ts2,
            "Graphics: Intel Ironlake": self.iron_gpu,
            "Graphics: Intel Sandy Bridge": self.sandy_gpu,
            "Graphics: Intel Ivy Bridge": self.ivy_gpu,
            # "Graphics: Intel Ivy Bridge": True,
            "Brightness: Legacy Backlight Control": self.brightness_legacy,
            "Audio: Legacy Realtek": self.legacy_audio,
            "Networking: Legacy Wireless": self.legacy_wifi,
            "Miscellaneous: Legacy GMUX": self.legacy_gmux,
            "Miscellaneous: Legacy Keyboard Backlight": self.legacy_keyboard_backlight,
            "Settings: Requires AMFI exemption": self.amfi_must_disable,
            "Validation: Patching Possible": self.verify_patch_allowed(),
            "Validation: SIP is enabled": self.sip_enabled,
            "Validation: SBM is enabled": self.sbm_enabled,
            "Validation: AMFI is enabled": self.amfi_enabled if self.amfi_must_disable else False,
            "Validation: FileVault is enabled": self.fv_enabled,
            "Validation: System is dosdude1 patched": self.dosdude_patched,
        }
        
        return self.root_patch_dict
    
    def verify_patch_allowed(self, print_errors=False):
        sip = sip_data.system_integrity_protection.root_patch_sip_big_sur if self.constants.detected_os > os_data.os_data.catalina else sip_data.system_integrity_protection.root_patch_sip_mojave
        self.sip_enabled, self.sbm_enabled, self.amfi_enabled, self.fv_enabled, self.dosdude_patched = utilities.patching_status(sip, self.constants.detected_os)
        if sip == sip_data.system_integrity_protection.root_patch_sip_mojave:
            sip_value = "For Hackintoshes, please set csr-active-config to '03060000' (0x603)\nFor non-OpenCore Macs, please run 'csrutil disable' in RecoveryOS"
        else:
            sip_value = (
                "For Hackintoshes, please set csr-active-config to '02080000' (0x802)\nFor non-OpenCore Macs, please run 'csrutil disable' and \n'csrutil authenticated-root disable' in RecoveryOS"
            )
        if print_errors is True:
            if self.sip_enabled is True:
                print("\nCannot patch! Please disable System Integrity Protection (SIP).")
                print("Disable SIP in Patcher Settings and Rebuild OpenCore\n")
                print("Ensure the following bits are set for csr-active-config:")
                print("\n".join(sip))
                print(sip_value)

            if self.sbm_enabled is True:
                print("\nCannot patch! Please disable Apple Secure Boot.")
                print("Disable SecureBootModel in Patcher Settings and Rebuild OpenCore")
                print("For Hackintoshes, set SecureBootModel to Disabled")

            if self.fv_enabled is True:
                print("\nCannot patch! Please disable FileVault.")
                print("For OCLP Macs, please rebuild your config with 0.2.5 or newer")
                print("For others, Go to System Preferences -> Security and disable FileVault")

            if self.amfi_enabled is True and self.amfi_must_disable is True:
                print("\nCannot patch! Please disable AMFI.")
                print("For Hackintoshes, please add amfi_get_out_of_my_way=1 to boot-args")
            
            if self.dosdude_patched is True:
                print("\nCannot patch! Detected machine has already been patched by another patcher")
                print("Please ensure your install is either clean or patched with OpenCore Legacy Patcher")
            
        if any(
            [self.sip_enabled, self.sbm_enabled, self.fv_enabled, self.dosdude_patched, self.amfi_enabled if self.amfi_must_disable else False]
        ):
            return False
        else:
            return True
    
    def generate_patchset(self, hardware_details):
        all_hardware_patchset = sys_patch_dict.SystemPatchDictionary(self.constants.detected_os)
        required_patches = {}
        utilities.cls()
        print("- The following patches will be applied:")
        if hardware_details["Graphics: Intel Ironlake"] is True:
            print("  - Graphics: Intel Ironlake")
            required_patches.update({"Non-Metal Common": all_hardware_patchset["Graphics"]["Non-Metal Common"]})
            required_patches.update({"Intel Ironlake": all_hardware_patchset["Graphics"]["Intel Ironlake"]})
        if hardware_details["Graphics: Intel Sandy Bridge"] is True:
            print("  - Graphics: Intel Sandy Bridge")
            required_patches.update({"Non-Metal Common": all_hardware_patchset["Graphics"]["Non-Metal Common"]})
            required_patches.update({"Legacy GVA": all_hardware_patchset["Graphics"]["Legacy GVA"]})
            required_patches.update({"Intel Sandy Bridge": all_hardware_patchset["Graphics"]["Intel Sandy Bridge"]})
        if hardware_details["Graphics: Intel Ivy Bridge"] is True:
            print("  - Graphics: Intel Ivy Bridge")
            required_patches.update({"Metal Common": all_hardware_patchset["Graphics"]["Metal Common"]})
            required_patches.update({"Intel Ivy Bridge": all_hardware_patchset["Graphics"]["Intel Ivy Bridge"]})
        if hardware_details["Graphics: Nvidia Tesla"] is True:
            print("  - Graphics: Nvidia Tesla")
            required_patches.update({"Non-Metal Common": all_hardware_patchset["Graphics"]["Non-Metal Common"]})
            required_patches.update({"Nvidia Tesla": all_hardware_patchset["Graphics"]["Nvidia Tesla"]})
            required_patches.update({"Nvidia Web Drivers": all_hardware_patchset["Graphics"]["Nvidia Web Drivers"]})
        if hardware_details["Graphics: Nvidia Kepler"] is True:
            print("  - Graphics: Nvidia Kepler")
            required_patches.update({"Metal Common": all_hardware_patchset["Graphics"]["Metal Common"]})
            required_patches.update({"Nvidia Kepler": all_hardware_patchset["Graphics"]["Nvidia Kepler"]})
        if hardware_details["Graphics: AMD TeraScale 1"] is True:
            print("  - Graphics: AMD TeraScale 1")
            required_patches.update({"Non-Metal Common": all_hardware_patchset["Graphics"]["Non-Metal Common"]})
            required_patches.update({"AMD Non-Metal Common": all_hardware_patchset["Graphics"]["AMD Non-Metal Common"]})
            required_patches.update({"AMD TeraScale 1": all_hardware_patchset["Graphics"]["AMD TeraScale 1"]})
        if hardware_details["Graphics: AMD TeraScale 2"] is True:
            print("  - Graphics: AMD TeraScale 2")
            required_patches.update({"Non-Metal Common": all_hardware_patchset["Graphics"]["Non-Metal Common"]})
            required_patches.update({"AMD Non-Metal Common": all_hardware_patchset["Graphics"]["AMD Non-Metal Common"]})
            required_patches.update({"AMD TeraScale 2": all_hardware_patchset["Graphics"]["AMD TeraScale 2"]})
        if hardware_details["Brightness: Legacy Backlight Control"] is True:
            print("  - Brightness: Legacy Brightness")
            required_patches.update({"Legacy Brightness": all_hardware_patchset["Brightness"]["Legacy Brightness"]})
        if hardware_details["Audio: Legacy Realtek"] is True:
            if self.model in ["iMac7,1", "iMac8,1"]:
                print("  - Audio: Legacy Realtek Audio")
                required_patches.update({"Legacy Realtek": all_hardware_patchset["Audio"]["Legacy Realtek"]})
            else:
                print("  - Audio: Legacy non-GOP Audio")
                required_patches.update({"Legacy Non-GOP": all_hardware_patchset["Audio"]["Legacy Non-GOP"]})
        if hardware_details["Networking: Legacy Wireless"] is True:
            print("  - Networking: Legacy WiFi")
            required_patches.update({"Legacy WiFi": all_hardware_patchset["Networking"]["Legacy WiFi"]})
        if hardware_details["Miscellaneous: Legacy GMUX"] is True:
            print("  - Miscellaneous: Legacy GMUX")
            required_patches.update({"Legacy GMUX": all_hardware_patchset["Miscellaneous"]["Legacy GMUX"]})
        if hardware_details["Miscellaneous: Legacy Keyboard Backlight"]:
            print("  - Miscellaneous: Legacy Keyboard Backlight")
            required_patches.update({"Legacy Keyboard Backlight": all_hardware_patchset["Miscellaneous"]["Legacy Keyboard Backlight"]})

        if not required_patches:
            print("  - No patch sets found for booted model")
        
        return required_patches