[Console]::OutputEncoding = [Text.Encoding]::UTF8
$ErrorActionPreference = 'Continue'

param(
    [string]$TargetUser = 'User'
)

Write-Host "[*] Adding desktop icons for user: $TargetUser"

# Icon GUIDs
$computerGUID = "{20D04FE0-3AEA-1069-A2D8-08002B30309D}"      # This PC / Computer
$controlPanelGUID = "{5399E694-6CE5-4D6C-8FCE-1D8870FDCBA0}"  # Control Panel
$recycleGUID = "{645FF040-5081-101B-9F08-00AA002F954E}"       # Recycle Bin

# Check if user exists
$userProfile = "C:\Users\$TargetUser"
if (-not (Test-Path $userProfile)) {
    Write-Host "[!] User profile not found: $userProfile" -ForegroundColor Red
    exit 1
}

Write-Host "[*] Target user profile: $userProfile"

# Check if user is currently logged in (hive already loaded)
$hivePath = "HKU:\${TargetUser}_Classes"
$userSID = $null

try {
    # Get user SID
    $user = New-Object System.Security.Principal.NTAccount($TargetUser)
    $userSID = $user.Translate([System.Security.Principal.SecurityIdentifier]).Value
    Write-Host "[*] User SID: $userSID"
} catch {
    Write-Host "[!] Could not get user SID: $($_.Exception.Message)" -ForegroundColor Red
    exit 1
}

# Check if user hive is already loaded
$hiveLoaded = $false
$userHivePath = "HKU:\$userSID"

if (Test-Path $userHivePath) {
    Write-Host "[*] User hive already loaded"
    $hiveLoaded = $true
} else {
    Write-Host "[*] Loading user hive"
    $ntUserPath = "$userProfile\NTUSER.DAT"

    if (-not (Test-Path $ntUserPath)) {
        Write-Host "[!] NTUSER.DAT not found: $ntUserPath" -ForegroundColor Red
        exit 1
    }

    try {
        & reg load "HKU\$TargetUser" "$ntUserPath"
        if ($LASTEXITCODE -eq 0) {
            Write-Host "    User hive loaded successfully"
            $userHivePath = "HKU:\$TargetUser"
        } else {
            Write-Host "[!] Failed to load user hive" -ForegroundColor Red
            exit 1
        }
    } catch {
        Write-Host "[!] Error loading user hive: $($_.Exception.Message)" -ForegroundColor Red
        exit 1
    }
}

# Desktop icons registry paths
$desktopIconsPath = "$userHivePath\Software\Microsoft\Windows\CurrentVersion\Explorer\HideDesktopIcons\NewStartPanel"
$classicIconsPath = "$userHivePath\Software\Microsoft\Windows\CurrentVersion\Explorer\HideDesktopIcons\ClassicStartMenu"

Write-Host "[*] Setting up desktop icons"

# Ensure registry paths exist
foreach ($path in @($desktopIconsPath, $classicIconsPath)) {
    if (-not (Test-Path $path)) {
        New-Item -Path $path -Force | Out-Null
        Write-Host "    Created registry path: $path"
    }
}

# Enable icons (0 = show, 1 = hide)
$iconsToEnable = @{
    "Computer (This PC)" = $computerGUID
    "Control Panel" = $controlPanelGUID
    "Recycle Bin" = $recycleGUID
}

foreach ($iconName in $iconsToEnable.Keys) {
    $guid = $iconsToEnable[$iconName]

    # Set for both NewStartPanel and ClassicStartMenu
    Set-ItemProperty -Path $desktopIconsPath -Name $guid -Value 0 -Type DWord
    Set-ItemProperty -Path $classicIconsPath -Name $guid -Value 0 -Type DWord

    Write-Host "    Enabled: $iconName"
}

# Unload user hive if we loaded it
if (-not $hiveLoaded) {
    Write-Host "[*] Unloading user hive"
    try {
        & reg unload "HKU\$TargetUser"
        if ($LASTEXITCODE -eq 0) {
            Write-Host "    User hive unloaded successfully"
        } else {
            Write-Host "[!] Warning: Failed to unload user hive cleanly"
        }
    } catch {
        Write-Host "[!] Warning: Error unloading user hive: $($_.Exception.Message)"
    }
}

Write-Host ""
Write-Host "[OK] Desktop icons configured for user: $TargetUser"
Write-Host ""
Write-Host "The following icons will appear on $TargetUser's desktop:"
Write-Host "- This PC (Computer)"
Write-Host "- Control Panel"
Write-Host "- Recycle Bin"
Write-Host ""
Write-Host "Changes will take effect when $TargetUser logs in next time."
Write-Host "If $TargetUser is currently logged in, they may need to:"
Write-Host "1. Log out and log back in"
Write-Host "2. Or restart the computer"
Write-Host "3. Or right-click desktop > Refresh"