# DeployCenter VNC module
# Run only with -Vnc in the main scope to share the WinPE context.

function Start-DeployCenterVnc {
    $parentRoot = Split-Path -Path $script:DeployRoot -Parent
    $sourceVnc = Join-Path $parentRoot 'uvnc'
    $sourceExe = Join-Path $sourceVnc 'winvnc.exe'
    $passwordTool = Join-Path $sourceVnc 'createpassword64.exe'
    $sourceIni = Join-Path $sourceVnc 'ultravnc.ini'

    if ([string]::IsNullOrWhiteSpace($script:DeployIPv4)) {
        & $script:SourceLogAction 'No deployment IPv4 address detected.' 'WARNING'
    }

    if (-not (Test-Path -LiteralPath $sourceExe -PathType Leaf)) {
        & $script:SourceLogAction "UltraVNC was not started: executable not found ($sourceExe)." 'WARNING'
        Write-DeployInventory -InventoryDir $script:InventoryDir -InventoryFile $script:InventoryFile -HardwareInfo $script:HardwareInventory -DeploymentIPv4 $script:DeployIPv4 -VncPassword $script:VncPassword -LogAction $script:SourceLogAction
        return
    }
    if (-not (Test-Path -LiteralPath $passwordTool -PathType Leaf)) {
        & $script:SourceLogAction "UltraVNC was not started: createpassword64.exe absent ($passwordTool)." 'WARNING'
        Write-DeployInventory -InventoryDir $script:InventoryDir -InventoryFile $script:InventoryFile -HardwareInfo $script:HardwareInventory -DeploymentIPv4 $script:DeployIPv4 -VncPassword $script:VncPassword -LogAction $script:SourceLogAction
        return
    }
    if (-not (Test-Path -LiteralPath $sourceIni -PathType Leaf)) {
        & $script:SourceLogAction "UltraVNC was not started: file not found ($sourceIni)." 'WARNING'
        Write-DeployInventory -InventoryDir $script:InventoryDir -InventoryFile $script:InventoryFile -HardwareInfo $script:HardwareInventory -DeploymentIPv4 $script:DeployIPv4 -VncPassword $script:VncPassword -LogAction $script:SourceLogAction
        return
    }

    $vncStage = 'preparing the temporary folder'
    try {
        if (Test-Path -LiteralPath $script:VncDirectory) {
            Remove-Item -LiteralPath $script:VncDirectory -Recurse -Force -ErrorAction Stop
        }
        New-Item -Path $script:VncDirectory -ItemType Directory -Force | Out-Null
        $vncStage = "copying files from $sourceVnc to $script:VncDirectory"
        Copy-Item -Path (Join-Path $sourceVnc '*') -Destination $script:VncDirectory -Recurse -Force
        # WinPE does not always allow firewall rules to be added.
        # The firewall is disabled globally only when VNC is requested.
        $vncStage = 'disabling the WinPE firewall'
        $wpeutil = Join-Path $env:WINDIR 'System32\wpeutil.exe'
        if (-not (Test-Path -LiteralPath $wpeutil -PathType Leaf)) {
            $wpeutil = 'wpeutil.exe'
        }
        $wpeutilOutput = @(& $wpeutil DisableFirewall 2>&1)
        $wpeutilExitCode = $LASTEXITCODE
        foreach ($outputLine in $wpeutilOutput) {
            $outputText = [string]$outputLine
            if (-not [string]::IsNullOrWhiteSpace($outputText)) {
                & $script:SourceLogAction "wpeutil DisableFirewall: $outputText" 'INFO' $true
            }
        }
        if ($wpeutilExitCode -ne 0) {
            throw "wpeutil DisableFirewall returned exit code $wpeutilExitCode."
        }
        & $script:SourceLogAction "WinPE firewall disabled at the request of -Vnc (wpeutil exit code: $wpeutilExitCode)." 'WARNING'

        # Force portable mode: the active INI is the one next to winvnc.exe.
        $vncStage = 'enabling UltraVNC portable mode'
        $portableFlag = Join-Path $script:VncDirectory 'ultravnc.portable'
        Set-Content -LiteralPath $portableFlag -Value '' -Encoding ASCII

        if (-not [string]::IsNullOrWhiteSpace([string]$script:Config.VncPassword)) {
            $script:VncPassword = [string]$script:Config.VncPassword
            $passwordSource = 'configured in DeployCenter.ini'
        }
        else {
            $script:VncPassword = -join (1..8 | ForEach-Object {
                $alphabet = 'ABCDEFGHJKLMNPQRSTUVWXYZabcdefghijkmnopqrstuvwxyz23456789'
                $alphabet[(Get-Random -Minimum 0 -Maximum $alphabet.Length)]
            })
            $passwordSource = 'generated automatically'
        }

        # createpassword64 encodes the selected password in the passwd field of ultravnc.ini.
        $vncStage = 'generating the UltraVNC password'
        $tool = Join-Path $script:VncDirectory 'createpassword64.exe'
        $script:VncProcess = Start-Process -FilePath $tool -ArgumentList @($script:VncPassword) `
            -WorkingDirectory $script:VncDirectory -Wait -PassThru -WindowStyle Hidden
        if ($script:VncProcess.ExitCode -ne 0) {
            throw "createpassword64.exe returned exit code $($script:VncProcess.ExitCode)."
        }

        $activeIni = Join-Path $script:VncDirectory 'ultravnc.ini'
        if (-not (Test-Path -LiteralPath $activeIni -PathType Leaf) -or
            -not (Select-String -LiteralPath $activeIni -Pattern '^passwd\s*=' -Quiet)) {
            throw "createpassword64.exe did not create the passwd field in ultravnc.ini."
        }

        $server = Join-Path $script:VncDirectory 'winvnc.exe'
        $vncStage = 'starting the UltraVNC server'
        $script:VncProcess = Start-Process -FilePath $server -ArgumentList @('-run') `
            -WorkingDirectory $script:VncDirectory -PassThru -WindowStyle Hidden
        Start-Sleep -Milliseconds 800
        if ($script:VncProcess.HasExited) {
            throw "winvnc.exe exited immediately (code $($script:VncProcess.ExitCode))."
        }
        $script:VncStarted = $true
        & $script:SourceLogAction "UltraVNC started. IP: $($script:DeployIPv4) ; password $passwordSource." 'SUCCESS'
    } catch {
        & $script:SourceLogAction "UltraVNC startup failed at stage '$vncStage' : $($_.Exception.Message)" 'ERROR'
        $script:VncStarted = $false
        $script:VncPassword = $null
    }
    Write-DeployInventory -InventoryDir $script:InventoryDir -InventoryFile $script:InventoryFile -HardwareInfo $script:HardwareInventory -DeploymentIPv4 $script:DeployIPv4 -VncPassword $script:VncPassword -LogAction $script:SourceLogAction
}

function Stop-DeployCenterVnc {
    try {
        if ($null -ne $script:VncProcess) {
            $proc = Get-Process -Id $script:VncProcess.Id -ErrorAction SilentlyContinue
            if ($null -ne $proc) { Stop-Process -Id $proc.Id -Force -ErrorAction SilentlyContinue }
        }
        # If winvnc starts a separate server process, target only processes
        # whose executable path matches our temporary copy.
        foreach ($proc in (Get-CimInstance Win32_Process -Filter "Name='winvnc.exe'" -ErrorAction SilentlyContinue)) {
            if ($proc.ExecutablePath -and $proc.ExecutablePath.StartsWith($script:VncDirectory, [StringComparison]::OrdinalIgnoreCase)) {
                Stop-Process -Id $proc.ProcessId -Force -ErrorAction SilentlyContinue
            }
        }
    } catch {
        & $script:SourceLogAction "Stopping UltraVNC: $($_.Exception.Message)" 'WARNING'
    }
    # The winvnc process may briefly keep a handle open on its folder.
    # The VNC folder belongs to this PowerShell process (uvnc_<PID>); do not
    # recursively delete the shared InventoryDir folder. This avoids
    # blocking the next launch if a handle is still open.
    try {
        if (Test-Path -LiteralPath $script:InventoryFile -PathType Leaf) {
            Remove-Item -LiteralPath $script:InventoryFile -Force -ErrorAction SilentlyContinue
        }
        if (Test-Path -LiteralPath $script:VncDirectory -PathType Container) {
            Start-Sleep -Milliseconds 500
            Remove-Item -LiteralPath $script:VncDirectory -Recurse -Force -ErrorAction SilentlyContinue
        }
    } catch { }
}