param(
    [switch]$InstallVirtIO,
    [switch]$EnableAutomaticLogon,
    [switch]$DisableAutoUpdate,
    [switch]$CopyInstallers,
    [switch]$InstallOpenSSH,
    [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 {
    # Numbered scripts (NN-*.bat / NN-*.ps1: installers and tools) go to Desktop\Tools.
    $scripts = @(Get-ChildItem -Path $PSScriptRoot -File |
        Where-Object { $_.Name -match '^\d\d-.*\.(bat|ps1)$' })
    if ($scripts.Count -eq 0) {
        throw "No installer batch files were found next to $PSCommandPath."
    }

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

# Installs the OpenSSH Server from the Win32-OpenSSH zip on the answer media,
# so no Windows Update download is needed.
function Install-OpenSSHServer {
    $zip = Get-ChildItem -Path $PSScriptRoot -Filter 'OpenSSH-*.zip' -File |
        Select-Object -First 1
    if (-not $zip) {
        throw "No OpenSSH-*.zip was found next to $PSCommandPath."
    }

    $installDir = Join-Path $env:ProgramFiles 'OpenSSH'
    if (-not (Test-Path -LiteralPath (Join-Path $installDir 'install-sshd.ps1'))) {
        $temp = Join-Path $env:TEMP 'openssh'
        Remove-Item $temp -Recurse -Force -ErrorAction SilentlyContinue
        Expand-Archive -LiteralPath $zip.FullName -DestinationPath $temp
        Remove-Item $installDir -Recurse -Force -ErrorAction SilentlyContinue
        Move-Item (Get-ChildItem $temp -Directory | Select-Object -First 1).FullName $installDir
        Remove-Item $temp -Recurse -Force
    }

    & powershell.exe -NoProfile -ExecutionPolicy Bypass -File (Join-Path $installDir 'install-sshd.ps1') | Out-Host
    if ($LASTEXITCODE -ne 0) {
        throw "install-sshd.ps1 failed (exit code $LASTEXITCODE)."
    }

    Set-Service -Name sshd -StartupType Automatic
    Start-Service -Name sshd

    if (-not (Get-NetFirewallRule -Name 'OpenSSH-Server-In-TCP' -ErrorAction SilentlyContinue)) {
        New-NetFirewallRule -Name 'OpenSSH-Server-In-TCP' -DisplayName 'OpenSSH Server (sshd)' `
            -Direction Inbound -Action Allow -Protocol TCP -LocalPort 22 -Profile Private, Public | Out-Null
    }

    # build-seed.sh puts the host's public key on the media; without it, only
    # password login works.
    $key = Join-Path $PSScriptRoot 'authorized_keys'
    if (-not (Test-Path -LiteralPath $key)) {
        Write-Warning "No authorized_keys was found next to $PSCommandPath; skipping the administrator key."
        return
    }
    $keyFile = Join-Path $env:ProgramData 'ssh\administrators_authorized_keys'
    Copy-Item -LiteralPath $key -Destination $keyFile -Force
    & icacls.exe $keyFile /inheritance:r /grant '*S-1-5-32-544:F' /grant '*S-1-5-18:F' | Out-Null
    if ($LASTEXITCODE -ne 0) {
        throw "icacls failed (exit code $LASTEXITCODE)."
    }
}

$exitCode = 0
try {
    if ($InstallVirtIO) {
        Install-VirtIOGuestTools
    }
    if ($EnableAutomaticLogon) {
        Enable-AutomaticLogon
    }
    if ($DisableAutoUpdate) {
        Disable-AutomaticUpdate
    }
    if ($CopyInstallers) {
        Copy-InstallerScripts
    }
    if ($InstallOpenSSH) {
        Install-OpenSSHServer
    }
    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