# Below setting allows to run local scripts without digital signature
Set-ExecutionPolicy -ExecutionPolicy RemoteSigned -Scope Process -Force

pnputil /scan-devices

$script_root = Split-Path -Parent $MyInvocation.MyCommand.Path
$driver_inf_path = Join-Path -Path $script_root -ChildPath "..\drivers\x64\npu.inf"
$driver_extn_inf_path= Join-Path -Path $script_root -ChildPath "..\drivers\x64\npu_extension_internal.inf"

$global:use_extn_inf = $true

# @Description
#       Enumerate connected devices and query specific PCI device with a given device ID on the machine
function Get-InstanceId {
	param( [string]$device_id )

	$instance_id = ""
	$pnputil_output = pnputil /enum-devices /connected | Select-String -Pattern "pci" | Select-String -Pattern $device_id

	foreach ($line in $pnputil_output) {
		$tokens = $line -split ":"
		$instance_id = $tokens[1].Trim()
		$instance_id = $instance_id -replace " "
	}

	return $instance_id
}

# @Description
#       Query the supported Device IDs to identify Device ID of the machine
function Get-DeviceInstanceId {

	$device_instance_id = ""

	if ($device_instance_id = Get-InstanceId -device_id "DEV_6240") {
		Write-Host "Device Instance ID: $device_instance_id"
		$global:use_extn_inf=$false
		return $device_instance_id
	}
	else {
		$device_ids = "DEV_AD1D", "DEV_7D1D", "DEV_643E", "DEV_B03E", "DEV_FD3E", "DEV_D71D"

		foreach ($device_id in $device_ids) {
			$device_instance_id = Get-InstanceId -device_id $device_id

			if($device_instance_id -ne "") {
				Write-Host "Device Instance ID: $device_instance_id"
				return $device_instance_id
			}
		}
	}

	if($device_instance_id -eq "") {
		Write-Host "Device not found"
		exit 1
	}
}

# @Description
#       Query the current status of the device
function Get-DeviceStatus {

	$status = ""
	$device_instance_id = Get-DeviceInstanceId
	$pnputil_output = pnputil /enum-devices /instanceid "$device_instance_id" | Select-String -Pattern "Status:"

	foreach ( $line in $pnputil_output ) {
		$tokens = $line -split ":"
		$status = $tokens[1].Trim()
		$status = $status -replace " "
	}

	if( $status -eq "Problem" ) {
		Write-Host "Device status: Problem"
		exit 1
	}
	elseif ( $status -eq "Started" ) {
		Write-Host "Device status: Driver installed"
	}
	elseif ( $status -eq "Stopped" ) {
		Write-Host "Device status: Driver not installed"
	}
	elseif ( $status -eq "Disabled" ) {
		Write-Host "Device status: Disabled"
	}
	elseif ( $status -eq "Disconnected" ) {
		Write-Host "Device status: Disconnected"
	}
	else {
		Write-Host "Device status: Unknown"
	}

	$result = [PSCustomObject]@{
		status = $status
		device_instance_id = $device_instance_id
	}

	return $result
}

# @Description
#       Query the oem driver name of the device
function Get-DriverName {
	param ([string]$device_instance_id)

	$pnputil_output = pnputil /enum-devices /instanceid "$device_instance_id" | Select-String -Pattern "Driver Name"
	$names = @()

	foreach ($line in $pnputil_output) {
		$tokens = $line -split ":"
		$name = $tokens[1].Trim()
		$name = $name -replace " "
		$names += $name
	}

	return $names
}

# @Description
#       Clean up the dangling npu and npu_extension infs
function Clear-NpuInf {

	$enumDriversOutput = pnputil /enum-drivers

	$originalNamePattern = 'Original Name:\s+(npu.*)'
	$publishedNamePattern = 'Published Name:\s+(.+)'

	$matches = $enumDriversOutput | Select-String -Pattern $originalNamePattern -Context 1

	foreach ($match in $matches) {
		$previousLine = $match.Context.PreContext[0]
		$publishedName = ($previousLine | Select-String -Pattern $publishedNamePattern).Matches.Groups[1].Value

		Write-Host "Deleting: $publishedName"
		pnputil /delete-driver $publishedName /uninstall /force
	}

	return
}

# @Description
#       Update/Delete registry name passed as an argument based on the action type.
#       At both state separated and absolute NPU registry path
# @args
#       Action types:
#          Update: Update the given registry name with value. Check if the key exists, and create them if they don't. Otherwise modify the registry value
#          Delete: Delete the given registry and whole state separated registry
function Update-Regkey($action, $name, $value, $infregkey) {
    $npu_pnp_class = "ComputeAccelerator"
    $npu_inf_reg_key = $infregkey
    $intel_abs_key_path = "HKLM:\SOFTWARE\Intel"
    $npu_abs_key_path = "HKLM:\SOFTWARE\Intel\VPU"

    # Fetch the NPU adapter class GUID
    $npu_class_guid = Get-CimInstance -Namespace root/cimv2 -ClassName Win32_PNPEntity | Where-Object PNPClass -EQ $npu_pnp_class | Select-Object -ExpandProperty ClassGuid

    # Scan for state separated registry key set by the inf
    $npu_inf_reg_key_paths = (Get-ChildItem -Path 'HKLM:/SYSTEM/ControlSet001/Control/Class/*' -Recurse -ErrorAction SilentlyContinue | Where-Object { $_.Property -eq $npu_inf_reg_key }).PSPATH

    if (-not (($npu_class_guid) -and ($npu_inf_reg_key_paths))){
        Write-Host "`r`nError**: NPU driver not installed correctly"
        Write-Host "npu_class_guid: ", $npu_class_guid
        Write-Host "npu_inf_reg_key_paths: ", $npu_inf_reg_key_paths, "`r`n"
        return
    }

    $npu_inf_reg_key_paths = $npu_inf_reg_key_paths.split("::")

    if($action -eq "Update" ){
        # Update state separated registry path
        ForEach ( $npu_inf_reg_key_path in $npu_inf_reg_key_paths){
            if ($npu_inf_reg_key_path -match $npu_class_guid){
                $npu_inf_reg_key_path = $npu_inf_reg_key_path.replace("HKEY_LOCAL_MACHINE", "HKLM:")
                Write-Host "`r`nWrting registry at :", $npu_inf_reg_key_path, ", Registry Name: ", $name," Value: ", $value, "`r`n"
                Set-ItemProperty -Path $npu_inf_reg_key_path -Name $name -Value $value -Type "DWord"
            }
        }

        # Update absolute NPU registry path, check if each key exists, and create them if they don't
        if (-not (Test-Path $intel_abs_key_path)) {
            Write-Host "Creating registry key: $intel_abs_key_path"
            New-Item -Path $intel_abs_key_path -Force
        }
        if (-not (Test-Path $npu_abs_key_path)) {
            Write-Host "Creating registry key: $npu_abs_key_path"
            New-Item -Path $npu_abs_key_path -Force
        }
        Write-Host "`r`nWrting registry at :", $npu_abs_key_path, ", Registry Name: ", $name," Value: ", $value, "`r`n"
        Set-ItemProperty -Path $npu_abs_key_path -Name $name -Value $value -Type "DWord"
    }
    elseif($action -eq "Delete"){
        # Delete state separated registry
        ForEach ( $npu_inf_reg_key_path in $npu_inf_reg_key_paths){
            if ($npu_inf_reg_key_path -match $npu_class_guid){
                $npu_inf_reg_key_path = $npu_inf_reg_key_path.replace("HKEY_LOCAL_MACHINE", "HKLM:")
                Write-Host "`r`nDeleting the driver store registry keys at: ", $npu_inf_reg_key_path, "`r`n"
                Remove-Item -Path $npu_inf_reg_key_path -Force -Recurse -ErrorAction SilentlyContinue
            }
        }

        # Delete absolute registry
        Remove-ItemProperty -Path $npu_abs_key_path -Name $name -Force -ErrorAction SilentlyContinue
    }
}

# @Description
#       Install the NPU driver
function Install-Driver {
	param ([string]$driver_inf_path)

	pnputil /add-driver $driver_inf_path /install | Tee-Object -Variable output
	Verify-PnputilResult $output

	return
}

# @Description
#       Install the npu_extension driver
function Install-ExtensionInf {
	param ([string]$driver_extn_inf_path)

	if( $global:use_extn_inf -ne $false) {
		Write-Host "Installing driver extension inf."
		pnputil /add-driver $driver_extn_inf_path /install | Tee-Object -Variable output
		Verify-PnputilResult $output
	}

	return
}

# @Description
#       Uninstall the driver
function Uninstall-Driver {
	param ([string]$device_instance_id, [string]$device_status)

	Get-DriverInfo -device_instance_id $device_instance_id

	$driver_names = Get-DriverName -device_instance_id $device_instance_id

	foreach ($driver_name in $driver_names) {
		if ( $driver_name -ne "" ) {
			pnputil /delete-driver $driver_name /uninstall /force | Tee-Object -Variable output
			Verify-PnputilResult $output
		}
	}

	Clear-NpuInf
	if ( $device_status -eq "Disconnected" ) {
		pnputil /remove-device /instanceid $device_instance_id
		pnputil /scan-devices
	}
	return
}

# @Description
#       Get the version and description of NPU driver
function Get-DriverInfo {
	param ([string]$device_instance_id)

	$device = Get-PnpDevice -PresentOnly | Where-Object { $_.InstanceId -eq $device_instance_id -and $_.FriendlyName -ne "PCI Device"  }

	if ($device) {
		Write-Host "Driver Description: $($device.FriendlyName)"
	}
	else {
		Write-Host "No driver found with DeviceID: $device_instance_id"
		return
	}

	$deviceDriver = Get-CimInstance Win32_PnPSignedDriver | Where-Object { $_.DeviceID -eq $device_instance_id }

	if ($deviceDriver) {
		Write-Host "Driver Version: $($deviceDriver.DriverVersion)"
	}

	return
}

# @Description
#       Disables a specific device on the system
function Enable-Device {
	param ([string]$device_instance_id)

	Write-Host "Enabling device"
	pnputil /enable-device /instanceid $device_instance_id | Tee-Object -Variable output
	Verify-PnputilResult $output

	return
}

# @Description
#       Disables a specific device on the system
function Disable-Device {
	param ([string]$device_instance_id)

	Write-Host "Disabling device"
	pnputil /disable-device /instanceid $device_instance_id | Tee-Object -Variable output
	Verify-PnputilResult $output

	return
}

# @Description
#       Verify the result of a pnputil command.
#       If the output contains an error message, exit with an error.
function Verify-PnputilResult {
	param ([string]$output)

	$errorMessages = @(
		"Failed to add driver package",
		"Failed to delete driver package",
		"Failed to enable device",
		"Failed to disable device",
		"Failed to scan devices"
	)

	foreach ($errorMessage in $errorMessages) {
		if ($output -match $errorMessage) {
			Write-Host "Error: $output"
			exit 1
		}
	}
}

$result = Get-DeviceStatus
$device_status = $result.status
$device_instance_id = $result.device_instance_id

if ( $args -contains "install_driver" ) {
	if ($args -notcontains "do_not_uninstall_existing_driver") {
		Uninstall-Driver -device_instance_id $device_instance_id -device_status $device_status
	}
	else {
		Write-Host "Skipping uninstallation of existing driver"
	}

	Install-Driver $driver_inf_path
	if ($args -notcontains "skip_extn_inf") {
		Install-ExtensionInf $driver_extn_inf_path
	}

	if ( $device_status -eq "Disabled" ) {
		Enable-Device -device_instance_id $device_instance_id
	}

	Get-DriverInfo -device_instance_id $device_instance_id

	if ($args -contains "reboot") {
		Write-Host "Rebooting the machine now..."
		Restart-Computer -Force
	}
}
elseif ( $args -contains "uninstall_driver" ) {
	Uninstall-Driver -device_instance_id $device_instance_id -device_status $device_status
}
elseif ( $args -contains "get_driver_info" ) {
	Get-DriverInfo -device_instance_id $device_instance_id
	exit 0
}
elseif ( $args -contains "enable_device" ) {
	if ( $device_status -ne "Enabled" )	{
		Enable-Device -device_instance_id $device_instance_id
	}
}
elseif ( $args -contains "disable_device" ) {
	if ( $device_status -ne "Disabled" ) {
		Disable-Device -device_instance_id $device_instance_id
	}
}
elseif ( $args -contains "clean" ) {
	Clear-NpuInf
}
else {
	Write-Host @"
	Invalid operation specified. Valid operations are:
	1.install_driver
		Flags that can be used with install_driver:
			reboot: Reboots the system after installation.
			do_not_uninstall_existing_driver: Keeps the current driver installed.
			skip_extn_inf: Installs base inf only
	2.install_driver reboot
	3.install_driver reboot do_not_uninstall_existing_driver
	4.install_driver do_not_uninstall_existing_driver
	5.uninstall_driver
	6.get_driver_info
	7.enable_device
	8.disable_device
	9.clean
"@
	exit 1
}

Get-DeviceStatus
