# ============================================================
# Set-WindowsDns.ps1
# ============================================================
# Remotely changes DNS settings on a Windows server.
# Connection order: WinRM (Invoke-Command) -> PSSession -> PsExec
# Includes DNS validation after the change.
# ============================================================

function Set-WindowsDns {
    <#
    .SYNOPSIS
        Changes DNS server addresses on a remote Windows machine.
    .DESCRIPTION
        Attempts connection via WinRM, then PSSession, then PsExec.
        In DryRun mode, reads and reports current DNS without making changes.
        Returns a structured result object compatible with Write-MigrationReport.
    .PARAMETER Hostname
        The hostname or FQDN of the target server.
    .PARAMETER NewDNS
        Array of new DNS server IP addresses to apply.
    .PARAMETER Config
        Hashtable loaded from config.psd1 containing connection settings.
    .PARAMETER Credential
        Optional PSCredential for Windows remote connections. If not provided,
        implicit authentication (current logged-in user / Kerberos) is used.
    .PARAMETER DryRun
        If set, no changes are made. Returns current DNS state and what would change.
    #>
    [CmdletBinding()]
    param(
        [Parameter(Mandatory)]
        [string]$Hostname,

        [Parameter(Mandatory)]
        [string[]]$NewDNS,

        [Parameter(Mandatory)]
        [hashtable]$Config,

        [System.Management.Automation.PSCredential]
        [System.Management.Automation.Credential()]
        $Credential = [System.Management.Automation.PSCredential]::Empty,

        [switch]$DryRun
    )

    $timestamp  = Get-Date -Format "yyyy-MM-ddTHH:mm:ss"
    $resolvedIP = Resolve-HostnameToIP -Hostname $Hostname

    # Base result object — populated throughout the function
    $result = [PSCustomObject]@{
        Hostname         = $Hostname
        ResolvedIP       = $resolvedIP
        OS               = "Windows"
        Distro           = ""
        OldDNS           = ""
        NewDNS           = $NewDNS -join ";"
        ConnectionMethod = ""
        ChangeStatus     = "Pending"
        ValidationStatus = "NotRun"
        ModifiedFiles    = "N/A"
        ServiceRestarted = "N/A"
        ErrorMessage     = ""
        DryRun           = $DryRun.IsPresent
        Timestamp        = $timestamp
    }

    # Cannot proceed without a resolvable hostname
    if ($resolvedIP -eq "Unresolvable") {
        $result.ChangeStatus     = "Skipped"
        $result.ValidationStatus = "Skipped"
        $result.ErrorMessage     = "Hostname could not be resolved to an IP address"
        Write-Warning "[$Hostname] Cannot resolve hostname — skipping"
        return $result
    }

    # ── Remote scriptblock ───────────────────────────────────────────────────
    # This block runs inside the remote session (WinRM or PSSession).
    # Returns a hashtable that is serialized back to the caller.
    $remoteScriptBlock = {
        param([string[]]$NewDNS, [bool]$DryRun, [string]$Domain)

        $output = @{
            OldDNS      = @()
            NICName     = ""
            OSCaption   = ""
            ChangeStatus = "Pending"
            ErrorMessage = ""
        }

        try {
            # Get OS caption for the result
            $os = (Get-WmiObject -Class Win32_OperatingSystem -ErrorAction Stop).Caption
            $output.OSCaption = $os

            # Find the NIC that has a default gateway (primary NIC)
            $defaultRoute = Get-NetRoute -DestinationPrefix "0.0.0.0/0" -ErrorAction SilentlyContinue |
                            Sort-Object RouteMetric |
                            Select-Object -First 1

            if ($null -eq $defaultRoute) {
                # Fallback: pick the first connected adapter
                $nic = Get-NetAdapter -Physical | Where-Object { $_.Status -eq "Up" } | Select-Object -First 1
            }
            else {
                $nic = Get-NetAdapter -InterfaceIndex $defaultRoute.InterfaceIndex -ErrorAction SilentlyContinue
            }

            if ($null -eq $nic) {
                $output.ChangeStatus  = "Failed"
                $output.ErrorMessage  = "No active network adapter found"
                return $output
            }

            $output.NICName = $nic.Name

            # Capture current DNS settings
            $currentDNS = Get-DnsClientServerAddress -InterfaceAlias $nic.Name -AddressFamily IPv4 -ErrorAction Stop
            $output.OldDNS = $currentDNS.ServerAddresses

            if ($DryRun) {
                $output.ChangeStatus = "DryRun"
                return $output
            }

            # Apply new DNS
            Set-DnsClientServerAddress -InterfaceAlias $nic.Name -ServerAddresses $NewDNS -ErrorAction Stop

            # Flush DNS cache
            Clear-DnsClientCache -ErrorAction SilentlyContinue

            $output.ChangeStatus = "Success"
        }
        catch {
            $output.ChangeStatus  = "Failed"
            $output.ErrorMessage  = $_.Exception.Message
        }

        return $output
    }

    # ── PsExec scriptblock ───────────────────────────────────────────────────
    # Used when WinRM and PSSession both fail.
    # Builds a one-liner PowerShell command to run via PsExec.
    $psexecScript = {
        param([string]$Hostname, [string[]]$NewDNS, [bool]$DryRun, [string]$PsExecPath, [string[]]$PsExecCredArgs)

        $newDnsJoined = '"{0}"' -f ($NewDNS -join '","')
        $dnsArray     = "@($newDnsJoined)"

        if ($DryRun) {
            $cmd = 'Get-DnsClientServerAddress -AddressFamily IPv4 | Select-Object -ExpandProperty ServerAddresses | ConvertTo-Json -Compress'
        }
        else {
            $cmd = @"
`$nic = (Get-NetRoute -DestinationPrefix '0.0.0.0/0' | Sort-Object RouteMetric | Select-Object -First 1 | ForEach-Object { Get-NetAdapter -InterfaceIndex `$_.InterfaceIndex }).Name;
if (-not `$nic) { `$nic = (Get-NetAdapter -Physical | Where-Object Status -eq Up | Select-Object -First 1).Name };
`$old = (Get-DnsClientServerAddress -InterfaceAlias `$nic -AddressFamily IPv4).ServerAddresses -join ',';
Set-DnsClientServerAddress -InterfaceAlias `$nic -ServerAddresses $dnsArray;
Clear-DnsClientCache;
Write-Output "OLD=`$old|NIC=`$nic|STATUS=success"
"@
        }

        $encodedCmd = [Convert]::ToBase64String([Text.Encoding]::Unicode.GetBytes($cmd))
        $psexecArgs = @("\\$Hostname", "-accepteula", "-nobanner", "-h") +
                      $PsExecCredArgs +
                      @("powershell.exe", "-NonInteractive", "-EncodedCommand", $encodedCmd)

        try {
            $psexecOutput = & $PsExecPath @psexecArgs 2>&1
            return $psexecOutput -join "`n"
        }
        catch {
            return "PSEXEC_ERROR: $($_.Exception.Message)"
        }
    }

    # ── Attempt WinRM connection ─────────────────────────────────────────────
    $remoteResult = $null
    $connectionMethod = ""

    # Build credential splatting — only added to params when credential is provided
    $credParams = @{}
    $useCredential = ($Credential -ne [System.Management.Automation.PSCredential]::Empty -and $null -ne $Credential)
    if ($useCredential) {
        $credParams.Credential = $Credential
        Write-Verbose "[$Hostname] Using explicit credentials: $($Credential.UserName)"
    }
    else {
        Write-Verbose "[$Hostname] Using implicit authentication (current user)"
    }

    Write-Verbose "[$Hostname] Attempting WinRM connection..."
    try {
        $invokeParams = @{
            ComputerName = $Hostname
            ScriptBlock  = $remoteScriptBlock
            ArgumentList = $NewDNS, $DryRun.IsPresent, $Config.Domain
            ErrorAction  = "Stop"
        }
        if ($Config.WinRMTimeout) {
            $sessionOpt = New-PSSessionOption -OpenTimeout ($Config.WinRMTimeout * 1000) -OperationTimeout ($Config.WinRMTimeout * 1000)
            $invokeParams.SessionOption = $sessionOpt
        }
        if ($useCredential) { $invokeParams.Credential = $Credential }

        $remoteResult        = Invoke-Command @invokeParams
        $connectionMethod    = "WinRM"
        Write-Verbose "[$Hostname] WinRM connection successful"
    }
    catch {
        Write-Warning "[$Hostname] WinRM failed: $($_.Exception.Message)"

        # ── Fallback: PSSession ──────────────────────────────────────────────
        Write-Verbose "[$Hostname] Attempting PSSession fallback..."
        try {
            $sessionOpt    = New-PSSessionOption -OpenTimeout ($Config.WinRMTimeout * 1000) -OperationTimeout ($Config.WinRMTimeout * 1000)
            $sessionParams = @{
                ComputerName  = $Hostname
                SessionOption = $sessionOpt
                ErrorAction   = "Stop"
            }
            if ($useCredential) { $sessionParams.Credential = $Credential }

            $session      = New-PSSession @sessionParams
            $remoteResult = Invoke-Command -Session $session -ScriptBlock $remoteScriptBlock `
                                -ArgumentList $NewDNS, $DryRun.IsPresent, $Config.Domain
            Remove-PSSession $session -ErrorAction SilentlyContinue
            $connectionMethod = "PSSession"
            Write-Verbose "[$Hostname] PSSession connection successful"
        }
        catch {
            Write-Warning "[$Hostname] PSSession failed: $($_.Exception.Message)"

            # ── Fallback: PsExec ─────────────────────────────────────────────
            Write-Verbose "[$Hostname] Attempting PsExec fallback..."
            if (-not (Test-Path $Config.PsExecPath)) {
                $result.ChangeStatus     = "Failed"
                $result.ErrorMessage     = "All connection methods failed. PsExec not found at: $($Config.PsExecPath)"
                $result.ConnectionMethod = "None"
                Write-Warning "[$Hostname] PsExec not found at $($Config.PsExecPath)"
                return $result
            }

            try {
                # Build PsExec credential args — only added when explicit credentials provided
                $psexecCredArgs = @()
                if ($useCredential) {
                    $psexecCredArgs = @(
                        "-u", $Credential.UserName,
                        "-p", $Credential.GetNetworkCredential().Password
                    )
                }

                $psexecOutput = & $psexecScript -Hostname $Hostname -NewDNS $NewDNS `
                                    -DryRun $DryRun.IsPresent -PsExecPath $Config.PsExecPath `
                                    -PsExecCredArgs $psexecCredArgs

                if ($psexecOutput -match "PSEXEC_ERROR:") {
                    throw $psexecOutput
                }

                # Parse PsExec output
                $connectionMethod = "PsExec"
                $remoteResult = @{ ChangeStatus = "Pending"; OldDNS = @(); NICName = ""; OSCaption = ""; ErrorMessage = "" }

                if ($DryRun) {
                    try {
                        $dnsArray = $psexecOutput | ConvertFrom-Json
                        $remoteResult.OldDNS      = $dnsArray
                        $remoteResult.ChangeStatus = "DryRun"
                    }
                    catch {
                        $remoteResult.OldDNS      = @($psexecOutput.Trim())
                        $remoteResult.ChangeStatus = "DryRun"
                    }
                }
                else {
                    if ($psexecOutput -match "STATUS=success") {
                        $remoteResult.ChangeStatus = "Success"
                        if ($psexecOutput -match "OLD=([^|]+)") { $remoteResult.OldDNS = $Matches[1] -split "," }
                        if ($psexecOutput -match "NIC=([^|]+)")  { $remoteResult.NICName = $Matches[1] }
                    }
                    else {
                        $remoteResult.ChangeStatus = "Failed"
                        $remoteResult.ErrorMessage = "PsExec output: $psexecOutput"
                    }
                }
                Write-Verbose "[$Hostname] PsExec connection successful"
            }
            catch {
                $result.ChangeStatus     = "Failed"
                $result.ConnectionMethod = "None"
                $result.ErrorMessage     = "All connection methods failed. Last error: $($_.Exception.Message)"
                Write-Warning "[$Hostname] All connection methods failed"
                return $result
            }
        }
    }

    # ── Populate result from remote output ───────────────────────────────────
    $result.ConnectionMethod = $connectionMethod
    $result.OldDNS           = if ($remoteResult.OldDNS) { $remoteResult.OldDNS -join ";" } else { "Unknown" }
    $result.ChangeStatus     = $remoteResult.ChangeStatus
    $result.ErrorMessage     = $remoteResult.ErrorMessage

    # Populate OS/Distro from remote caption if available
    if ($remoteResult.NICName) {
        $result.ModifiedFiles = "NIC:$($remoteResult.NICName)"
    }
    if ($remoteResult.OSCaption) {
        $result.Distro = $remoteResult.OSCaption
    }

    # ── Validate DNS resolution (Task 3) ─────────────────────────────────────
    if ($result.ChangeStatus -in @("Success", "DryRun")) {
        $result.ValidationStatus = Test-WindowsDnsValidation `
            -Hostname         $Hostname `
            -NewDNS           $NewDNS[0] `
            -Domain           $Config.Domain `
            -ConnectionMethod $connectionMethod `
            -Config           $Config `
            -Credential       $Credential `
            -DryRun:          $DryRun
    }

    return $result
}


function Test-WindowsDnsValidation {
    <#
    .SYNOPSIS
        Validates that a remote Windows server can resolve the domain via the new DNS IP.
    .DESCRIPTION
        Runs Resolve-DnsName against the new DNS server from the remote machine.
        In DryRun mode, tests resolution against current DNS (non-destructive).
    #>
    [CmdletBinding()]
    param(
        [Parameter(Mandatory)] [string]$Hostname,
        [Parameter(Mandatory)] [string]$NewDNS,
        [Parameter(Mandatory)] [string]$Domain,
        [Parameter(Mandatory)] [string]$ConnectionMethod,
        [Parameter(Mandatory)] [hashtable]$Config,
        [System.Management.Automation.PSCredential]
        [System.Management.Automation.Credential()]
        $Credential = [System.Management.Automation.PSCredential]::Empty,
        [switch]$DryRun
    )

    $dnsToTest     = $NewDNS
    $useCredential = ($Credential -ne [System.Management.Automation.PSCredential]::Empty -and $null -ne $Credential)
    Write-Verbose "[$Hostname] Validating DNS resolution of '$Domain' via $dnsToTest..."

    $validationBlock = {
        param([string]$Domain, [string]$DnsServer)
        try {
            $result = Resolve-DnsName -Name $Domain -Server $DnsServer -Type A -ErrorAction Stop
            if ($result) {
                return "Success|Resolved $Domain to: $(($result | Where-Object { $_.Type -eq 'A' } | Select-Object -ExpandProperty IPAddress) -join ', ')"
            }
            return "ValidationFailed|No A records returned for $Domain"
        }
        catch {
            return "ValidationFailed|$($_.Exception.Message)"
        }
    }

    $validationOutput = $null

    try {
        switch ($ConnectionMethod) {
            "WinRM" {
                $sessionOpt    = New-PSSessionOption -OpenTimeout ($Config.WinRMTimeout * 1000)
                $invokeParams  = @{
                    ComputerName  = $Hostname
                    ScriptBlock   = $validationBlock
                    ArgumentList  = $Domain, $dnsToTest
                    SessionOption = $sessionOpt
                    ErrorAction   = "Stop"
                }
                if ($useCredential) { $invokeParams.Credential = $Credential }
                $validationOutput = Invoke-Command @invokeParams
            }
            "PSSession" {
                $sessionOpt    = New-PSSessionOption -OpenTimeout ($Config.WinRMTimeout * 1000)
                $sessionParams = @{ ComputerName = $Hostname; SessionOption = $sessionOpt; ErrorAction = "Stop" }
                if ($useCredential) { $sessionParams.Credential = $Credential }
                $session = New-PSSession @sessionParams
                $validationOutput = Invoke-Command -Session $session -ScriptBlock $validationBlock `
                    -ArgumentList $Domain, $dnsToTest
                Remove-PSSession $session -ErrorAction SilentlyContinue
            }
            "PsExec" {
                # Run nslookup via PsExec as Resolve-DnsName may not be available on older OS
                $cmd     = "nslookup $Domain $dnsToTest"
                $encoded = [Convert]::ToBase64String([Text.Encoding]::Unicode.GetBytes($cmd))
                $credArgs = @()
                if ($useCredential) {
                    $credArgs = @("-u", $Credential.UserName, "-p", $Credential.GetNetworkCredential().Password)
                }
                $psexecArgs = @("\\$Hostname", "-accepteula", "-nobanner", "-h") +
                              $credArgs +
                              @("powershell.exe", "-NonInteractive", "-EncodedCommand", $encoded)
                $output = & $Config.PsExecPath @psexecArgs 2>&1
                if ($LASTEXITCODE -eq 0 -and $output -match "Address") {
                    $validationOutput = "Success|nslookup succeeded: $($output -join ' ')"
                }
                else {
                    $validationOutput = "ValidationFailed|nslookup output: $($output -join ' ')"
                }
            }
        }
    }
    catch {
        return "ValidationFailed|Validation connection error: $($_.Exception.Message)"
    }

    if ($null -eq $validationOutput) {
        return "ValidationFailed|No output from validation check"
    }

    $parts = $validationOutput -split "\|", 2
    Write-Verbose "[$Hostname] Validation result: $($parts[0])"
    return $parts[0]
}


function Resolve-HostnameToIP {
    <#
    .SYNOPSIS
        Resolves a hostname to its first IPv4 address. Returns "Unresolvable" on failure.
    #>
    param([string]$Hostname)

    # If already an IP, return as-is
    if ($Hostname -match '^\d{1,3}(\.\d{1,3}){3}$') { return $Hostname }

    try {
        $addresses = [System.Net.Dns]::GetHostAddresses($Hostname) |
                     Where-Object { $_.AddressFamily -eq 'InterNetwork' }
        if ($addresses) {
            return ($addresses | Select-Object -First 1).IPAddressToString
        }
        return "Unresolvable"
    }
    catch {
        return "Unresolvable"
    }
}