param(
    [switch]$InstallVirtIO,
    [switch]$EnableAutomaticLogon,
    [switch]$DisableAutoUpdate,
    [switch]$CopyInstallers,
    [switch]$NoPause
)

$ErrorActionPreference = 'Stop'
Set-StrictMode -Version Latest

function Test-Administrator {
    $identity = [Security.Principal.WindowsIdentity]::GetCurrent()
    $principal = [Security.Principal.WindowsPrincipal]::new($identity)
    return $principal.IsInRole([Security.Principal.WindowsBuiltInRole]::Administrator)
}

if (-not (Test-Administrator)) {
    if (-not $PSCommandPath) {
        throw 'Run this script from a PowerShell window opened as Administrator.'
    }

    $arguments = @('-NoProfile', '-ExecutionPolicy', 'Bypass', '-File', "`"$PSCommandPath`"")
    foreach ($name in $PSBoundParameters.Keys) {
        $arguments += "-$name"
    }

    $elevated = Start-Process `
        -FilePath 'powershell.exe' `
        -ArgumentList $arguments `
        -Verb RunAs `
        -Wait `
        -PassThru
    exit $elevated.ExitCode
}

# The guest tools installer is an x86 program and the VirtIO ISO has no Arm64
# package for it, so on Arm64 Windows the Arm64 drivers are installed from their
# folders on the ISO instead. NetKVM is the network driver; without it the guest
# has no network.
function Install-VirtIOArm64Drivers {
    $roots = (Get-PSDrive -PSProvider FileSystem).Root
    $drivers = @(
        @{ Folder = 'NetKVM'; Inf = 'netkvm.inf' },
        @{ Folder = 'vioserial'; Inf = 'vioser.inf' },
        @{ Folder = 'Balloon'; Inf = 'balloon.inf' },
        @{ Folder = 'viorng'; Inf = 'viorng.inf' },
        @{ Folder = 'vioinput'; Inf = 'vioinput.inf' },
        @{ Folder = 'viogpudo'; Inf = 'viogpudo.inf' }
    )

    $installed = 0
    foreach ($driver in $drivers) {
        $inf = $roots |
            ForEach-Object { Join-Path $_ "$($driver.Folder)\w11\ARM64\$($driver.Inf)" } |
            Where-Object { Test-Path -LiteralPath $_ } |
            Select-Object -First 1
        if (-not $inf) {
            if ($driver.Folder -eq 'NetKVM') {
                throw 'The Arm64 NetKVM driver was not found on the attached media.'
            }
            continue
        }

        & pnputil.exe /add-driver $inf /install | Out-Host
        # 3010 asks for a restart; 259 means that no device uses the driver.
        if ($LASTEXITCODE -notin 0, 259, 3010) {
            throw "pnputil failed to install $inf (exit code $LASTEXITCODE)."
        }
        $installed++
    }
    Write-Host "Installed $installed VirtIO Arm64 driver packages."
}

function Install-VirtIOGuestTools {
    if ([Runtime.InteropServices.RuntimeInformation]::OSArchitecture -eq 'Arm64') {
        Install-VirtIOArm64Drivers
        return
    }

    $installer = Get-ChildItem `
        -Path (Get-PSDrive -PSProvider FileSystem).Root `
        -Filter 'virtio-win-guest-tools.exe' `
        -File `
        -ErrorAction SilentlyContinue |
        Select-Object -First 1

    if (-not $installer) {
        throw 'virtio-win-guest-tools.exe was not found on the attached media.'
    }

    $process = Start-Process `
        -FilePath $installer.FullName `
        -ArgumentList '/quiet', '/norestart' `
        -Wait `
        -PassThru

    if ($process.ExitCode -notin 0, 3010) {
        throw "VirtIO guest tools installer exited with code $($process.ExitCode)."
    }
}

function Enable-AutomaticLogon {
    $winlogonPath = 'HKLM:\SOFTWARE\Microsoft\Windows NT\CurrentVersion\Winlogon'

    Set-ItemProperty -Path $winlogonPath -Name 'AutoAdminLogon' -Value '1' -Type String
    Set-ItemProperty -Path $winlogonPath -Name 'DefaultUserName' -Value 'guofang' -Type String
    Set-ItemProperty -Path $winlogonPath -Name 'DefaultPassword' -Value 'wang' -Type String
    Set-ItemProperty -Path $winlogonPath -Name 'DefaultDomainName' -Value $env:COMPUTERNAME -Type String
    Remove-ItemProperty -Path $winlogonPath -Name 'AutoLogonCount' -ErrorAction SilentlyContinue
}

function Disable-AutomaticUpdate {
    $windowsUpdatePath = 'HKLM:\SOFTWARE\Policies\Microsoft\Windows\WindowsUpdate'
    $autoUpdatePath = Join-Path $windowsUpdatePath 'AU'
    $storePath = 'HKLM:\SOFTWARE\Policies\Microsoft\WindowsStore'
    $currentVersion = Get-ItemProperty `
        -Path 'HKLM:\SOFTWARE\Microsoft\Windows NT\CurrentVersion'
    $releaseVersion = $currentVersion.DisplayVersion

    New-Item -Path $autoUpdatePath -Force | Out-Null
    New-Item -Path $storePath -Force | Out-Null

    # Never check automatically; manual checks and installs keep working.
    Set-ItemProperty -Path $autoUpdatePath -Name 'NoAutoUpdate' -Value 1 -Type DWord
    Set-ItemProperty -Path $autoUpdatePath -Name 'NoAutoRebootWithLoggedOnUsers' -Value 1 -Type DWord
    Set-ItemProperty -Path $windowsUpdatePath -Name 'ExcludeWUDriversInQualityUpdate' -Value 1 -Type DWord

    # Stay on the installed feature release instead of upgrading automatically.
    if ($releaseVersion) {
        Set-ItemProperty -Path $windowsUpdatePath -Name 'TargetReleaseVersion' -Value 1 -Type DWord
        Set-ItemProperty -Path $windowsUpdatePath -Name 'TargetReleaseVersionInfo' -Value $releaseVersion -Type String
    }

    # 2 = never download Microsoft Store app updates automatically.
    Set-ItemProperty -Path $storePath -Name 'AutoDownload' -Value 2 -Type DWord

    # Windows may re-enable this task, so the policies above are the real switch.
    Disable-ScheduledTask `
        -TaskPath '\Microsoft\Windows\UpdateOrchestrator\' `
        -TaskName 'Schedule Scan' `
        -ErrorAction SilentlyContinue | Out-Null
}

function Copy-InstallerScripts {
    $scripts = @(Get-ChildItem -Path $PSScriptRoot -Filter '*-Install-*.bat' -File)
    if ($scripts.Count -eq 0) {
        throw "No installer batch files were found next to $PSCommandPath."
    }

    $target = Join-Path ([Environment]::GetFolderPath('Desktop')) 'Installers'
    New-Item -ItemType Directory -Path $target -Force | Out-Null
    Copy-Item -Path $scripts.FullName -Destination $target -Force

    # 05-Install-OpenSSH-Admin.bat installs from this zip when it is present.
    $zips = @(Get-ChildItem -Path $PSScriptRoot -Filter 'OpenSSH-*.zip' -File)
    if ($zips.Count -gt 0) {
        Copy-Item -Path $zips.FullName -Destination $target -Force
    }
}

$exitCode = 0
try {
    if ($InstallVirtIO) {
        Install-VirtIOGuestTools
    }
    if ($EnableAutomaticLogon) {
        Enable-AutomaticLogon
    }
    if ($DisableAutoUpdate) {
        Disable-AutomaticUpdate
    }
    if ($CopyInstallers) {
        Copy-InstallerScripts
    }
    Write-Host 'Done.' -ForegroundColor Green
}
catch {
    Write-Host "Failed: $($_.Exception.Message)" -ForegroundColor Red
    $exitCode = 1
}

if (-not $NoPause) {
    Read-Host 'Press Enter to exit' | Out-Null
}
exit $exitCode