#Requires -Version 5.1
<#
.SYNOPSIS
    Restores DNS settings on Windows and Linux servers from a rollback JSON file.

.DESCRIPTION
    Reads a rollback JSON file produced by Invoke-DnsMigration.ps1 and
    restores each server's DNS configuration to its pre-migration state.

    For Windows servers:
      Reconnects via WinRM/PSSession/PsExec and calls
      Set-DnsClientServerAddress with the original DNS IPs.

    For Linux servers:
      SSHes in via plink, overwrites each modified config file with its
      original content (stored as base64 in the rollback JSON), then
      restarts the appropriate service.

    Supports selective rollback (-Hostname) and dry-run mode (-DryRun).

.PARAMETER RollbackFile
    Path to the dns-rollback-<timestamp>.json file generated during migration.
    Mandatory.

.PARAMETER Hostname
    Optional. Limit rollback to a single server. Case-insensitive match
    against the Hostname field in the rollback file.

.PARAMETER ConfigFile
    Path to config.psd1. Defaults to .\config.psd1 in the script directory.

.PARAMETER OutputDir
    Directory for the rollback CSV report. Defaults to config.OutputDir.

.PARAMETER LinuxSSHPassword
    Override the LinuxSSHPassword from config.psd1 at runtime.

.PARAMETER WindowsCredential
    Optional PSCredential for Windows remote connections.
    If not provided, uses implicit authentication (current user / Kerberos).
    Example: -WindowsCredential (Get-Credential)

.PARAMETER WindowsUsername
    Optional Windows username. Must be used with -WindowsPassword.
    Example: -WindowsUsername "CORP\migration-svc" -WindowsPassword "P@ssw0rd"

.PARAMETER DryRun
    Read-only mode. Shows what would be restored (first 10 lines of each
    file's original content) without making any changes.

.EXAMPLE
    # Dry run — preview full fleet rollback
    .\Invoke-DnsRollback.ps1 -RollbackFile ".\logs\dns-rollback-20250115-103000.json" -DryRun

.EXAMPLE
    # Roll back a single server
    .\Invoke-DnsRollback.ps1 -RollbackFile ".\logs\dns-rollback-20250115-103000.json" `
                              -Hostname "ubuntu-app-01.corp.local"

.EXAMPLE
    # Full fleet rollback
    .\Invoke-DnsRollback.ps1 -RollbackFile ".\logs\dns-rollback-20250115-103000.json"
#>
[CmdletBinding(SupportsShouldProcess)]
param(
    [Parameter(Mandatory)]
    [string] $RollbackFile,

    [string] $Hostname         = "",
    [string] $ConfigFile       = "",
    [string] $OutputDir        = "",
    [string] $LinuxSSHPassword = "",

    # ── Windows authentication (optional — omit for implicit Kerberos) ────────
    [System.Management.Automation.PSCredential]
    [System.Management.Automation.Credential()]
    $WindowsCredential = [System.Management.Automation.PSCredential]::Empty,

    [string] $WindowsUsername  = "",
    [string] $WindowsPassword  = "",

    [switch] $DryRun
)

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

$ScriptDir = if ($PSScriptRoot) { $PSScriptRoot } else { Split-Path $MyInvocation.MyCommand.Path }

# ── Load modules ──────────────────────────────────────────────────────────────
$modulePaths = @(
    "modules\Set-WindowsDns.ps1",
    "modules\Invoke-LinuxDns.ps1",
    "modules\Write-MigrationReport.ps1"
)
foreach ($mp in $modulePaths) {
    $fullPath = Join-Path $ScriptDir $mp
    if (-not (Test-Path $fullPath)) { throw "Required module not found: $fullPath" }
    . $fullPath
}

# ── Load config ───────────────────────────────────────────────────────────────
if (-not $ConfigFile) { $ConfigFile = Join-Path $ScriptDir "config.psd1" }
if (-not (Test-Path $ConfigFile)) { throw "Configuration file not found: $ConfigFile" }

$config = Import-PowerShellDataFile $ConfigFile

if ($LinuxSSHPassword) { $config.LinuxSSHPassword = $LinuxSSHPassword }
if ($OutputDir)        { $config.OutputDir        = $OutputDir }

# ── Resolve Windows credential ────────────────────────────────────────────────
$winCredential = [System.Management.Automation.PSCredential]::Empty

if ($WindowsCredential -ne [System.Management.Automation.PSCredential]::Empty -and $null -ne $WindowsCredential) {
    $winCredential = $WindowsCredential
}
elseif ($WindowsUsername -and $WindowsPassword) {
    $securePass    = ConvertTo-SecureString $WindowsPassword -AsPlainText -Force
    $winCredential = New-Object System.Management.Automation.PSCredential($WindowsUsername, $securePass)
}
elseif ($config.WindowsUsername -and $config.WindowsPassword) {
    $securePass    = ConvertTo-SecureString $config.WindowsPassword -AsPlainText -Force
    $winCredential = New-Object System.Management.Automation.PSCredential($config.WindowsUsername, $securePass)
}

$config._WinCredential = $winCredential

if (-not [System.IO.Path]::IsPathRooted($config.OutputDir)) {
    $config.OutputDir = Join-Path $ScriptDir $config.OutputDir
}
if (-not (Test-Path $config.OutputDir)) {
    New-Item -ItemType Directory -Path $config.OutputDir -Force | Out-Null
}

# ── Banner ────────────────────────────────────────────────────────────────────
$runMode = if ($DryRun) { "DRY RUN — no changes will be made" } else { "LIVE ROLLBACK — DNS settings WILL be restored" }
Write-Host ""
Write-Host "╔══════════════════════════════════════════════════════════════╗" -ForegroundColor Yellow
Write-Host "║           DNS ROLLBACK ORCHESTRATOR                         ║" -ForegroundColor Yellow
Write-Host "╚══════════════════════════════════════════════════════════════╝" -ForegroundColor Yellow
Write-Host "  Mode         : " -NoNewline
if ($DryRun) { Write-Host $runMode -ForegroundColor Yellow }
else         { Write-Host $runMode -ForegroundColor Red    }
Write-Host "  Rollback file: $RollbackFile"
if ($Hostname) {
    Write-Host "  Scope        : Single server — $Hostname" -ForegroundColor Cyan
}
else {
    Write-Host "  Scope        : Full fleet"
}
Write-Host "  Win auth     : " -NoNewline
if ($winCredential -ne [System.Management.Automation.PSCredential]::Empty) {
    Write-Host "Explicit ($($winCredential.UserName))" -ForegroundColor Yellow
} else {
    Write-Host "Implicit (current user / Kerberos)" -ForegroundColor Gray
}
Write-Host "  Started at   : $(Get-Date -Format 'yyyy-MM-dd HH:mm:ss')"
Write-Host ""

# ── Load rollback file ────────────────────────────────────────────────────────
Write-Host "── Loading rollback data ────────────────────────────────────────" -ForegroundColor Yellow
$rollbackData = Read-RollbackFile -Path $RollbackFile

# ── Filter entries ────────────────────────────────────────────────────────────
$entries = @($rollbackData.Servers)

if ($Hostname) {
    $entries = @($entries | Where-Object { $_.Hostname -ieq $Hostname })
    if ($entries.Count -eq 0) {
        Write-Warning "No entry found for hostname '$Hostname' in rollback file."
        Write-Host "Available hostnames:"
        $rollbackData.Servers | ForEach-Object { Write-Host "  $($_.Hostname)" }
        exit 1
    }
    Write-Host "  Filtered to 1 server: $Hostname" -ForegroundColor Cyan
}

Write-Host "  Processing $($entries.Count) server(s) for rollback" -ForegroundColor Yellow
Write-Host ""

# ── Process rollback entries ──────────────────────────────────────────────────
Write-Host "── Rolling back servers ─────────────────────────────────────────" -ForegroundColor Yellow

$allResults = [System.Collections.Generic.List[PSCustomObject]]::new()
$i = 0

foreach ($entry in $entries) {
    $i++
    $entryHostname = $entry.Hostname
    $entryOS       = $entry.OS

    Write-Host "  [$i/$($entries.Count)] $entryHostname ($entryOS) ..." -NoNewline

    $result = $null

    try {
        if ($entryOS -eq "Windows") {
            $result = Invoke-WindowsRollback -Entry $entry -Config $config -Credential $config._WinCredential -DryRun:$DryRun
        }
        elseif ($entryOS -eq "Linux") {
            # Convert PSCustomObject entry to hashtable for Invoke-LinuxDnsRollback
            $entryHT = ConvertTo-RollbackHashtable -Entry $entry
            $result  = Invoke-LinuxDnsRollback -RollbackEntry $entryHT -Config $config -DryRun:$DryRun
        }
        else {
            $result = [PSCustomObject]@{
                Hostname         = $entryHostname
                ResolvedIP       = "Unknown"
                OS               = $entryOS
                Distro           = $entry.Distro
                OldDNS           = ($entry.NewDNS -join ";")
                NewDNS           = ($entry.OldDNS -join ";")
                ConnectionMethod = "None"
                ChangeStatus     = "Skipped"
                ValidationStatus = "Skipped"
                ModifiedFiles    = ""
                ServiceRestarted = ""
                ErrorMessage     = "Unknown OS in rollback file: $entryOS"
                DryRun           = $DryRun.IsPresent
                Timestamp        = (Get-Date -Format "yyyy-MM-ddTHH:mm:ss")
            }
        }
    }
    catch {
        $result = [PSCustomObject]@{
            Hostname         = $entryHostname
            ResolvedIP       = "Unknown"
            OS               = $entryOS
            Distro           = $entry.Distro
            OldDNS           = ($entry.NewDNS -join ";")
            NewDNS           = ($entry.OldDNS -join ";")
            ConnectionMethod = "Error"
            ChangeStatus     = "Failed"
            ValidationStatus = "NotRun"
            ModifiedFiles    = ""
            ServiceRestarted = ""
            ErrorMessage     = "Rollback exception: $($_.Exception.Message)"
            DryRun           = $DryRun.IsPresent
            Timestamp        = (Get-Date -Format "yyyy-MM-ddTHH:mm:ss")
        }
    }

    $statusColor = switch ($result.ChangeStatus) {
        "Success" { "Green"  }
        "DryRun"  { "Cyan"   }
        "Failed"  { "Red"    }
        "Skipped" { "Yellow" }
        default   { "White"  }
    }
    Write-Host " [$($result.ChangeStatus)]" -ForegroundColor $statusColor -NoNewline
    Write-Host " (Validation: $($result.ValidationStatus))"

    $allResults.Add($result)
}

# ── Write rollback report ─────────────────────────────────────────────────────
$runTimestamp      = Get-Date -Format "yyyyMMdd-HHmmss"
$rollbackReportPath = Join-Path $config.OutputDir "dns-rollback-report-$runTimestamp.csv"

Write-Host ""
Write-Host "── Writing rollback report ──────────────────────────────────────" -ForegroundColor Yellow
Write-MigrationReport -Results $allResults.ToArray() -OutputPath $rollbackReportPath -NoAppend

Write-Host "  Rollback completed at: $(Get-Date -Format 'yyyy-MM-dd HH:mm:ss')" -ForegroundColor Gray
Write-Host ""


# ── Windows rollback function ─────────────────────────────────────────────────
function Invoke-WindowsRollback {
    <#
    .SYNOPSIS
        Restores original DNS settings on a Windows server from rollback data.
    .DESCRIPTION
        Connects via WinRM/PSSession/PsExec and calls Set-DnsClientServerAddress
        with the original DNS IPs stored in the rollback file.
    #>
    param(
        [object]   $Entry,
        [hashtable]$Config,
        [System.Management.Automation.PSCredential]
        [System.Management.Automation.Credential()]
        $Credential = [System.Management.Automation.PSCredential]::Empty,
        [switch]   $DryRun
    )

    $hostname      = $Entry.Hostname
    $oldDNS        = @($Entry.OldDNS)   # This is what we restore TO (pre-migration DNS)
    $nicName       = $Entry.NICName
    $timestamp     = Get-Date -Format "yyyy-MM-ddTHH:mm:ss"
    $useCredential = ($Credential -ne [System.Management.Automation.PSCredential]::Empty -and $null -ne $Credential)

    $resolvedIP = Resolve-ADHostnameToIP -Hostname $hostname

    $result = [PSCustomObject]@{
        Hostname         = $hostname
        ResolvedIP       = $resolvedIP
        OS               = "Windows"
        Distro           = $Entry.Distro
        OldDNS           = ($Entry.NewDNS -join ";")   # Rolling back FROM NewDNS
        NewDNS           = ($oldDNS -join ";")          # Rolling back TO OldDNS
        ConnectionMethod = ""
        ChangeStatus     = "Pending"
        ValidationStatus = "NotRun"
        ModifiedFiles    = if ($nicName) { "NIC:$nicName" } else { "N/A" }
        ServiceRestarted = "N/A"
        ErrorMessage     = ""
        DryRun           = $DryRun.IsPresent
        Timestamp        = $timestamp
    }

    if ($resolvedIP -eq "Unresolvable") {
        $result.ChangeStatus     = "Skipped"
        $result.ValidationStatus = "Skipped"
        $result.ErrorMessage     = "Hostname unresolvable"
        return $result
    }

    if ($DryRun) {
        Write-Host ""
        Write-Host "  [$hostname] ROLLBACK DRY RUN (Windows)" -ForegroundColor Yellow
        Write-Host "    NIC         : $nicName" -ForegroundColor Cyan
        Write-Host "    Restore DNS : $($oldDNS -join ', ')" -ForegroundColor Cyan
        Write-Host "    Current DNS : $($Entry.NewDNS -join ', ')" -ForegroundColor Gray
        $result.ChangeStatus     = "DryRun"
        $result.ValidationStatus = "DryRun"
        return $result
    }

    # Remote scriptblock to restore DNS
    $restoreBlock = {
        param([string[]]$OldDNS, [string]$NICName)
        $out = @{ Status="Pending"; Error="" }
        try {
            $targetNIC = $NICName
            if (-not $targetNIC -or $targetNIC -eq "") {
                # Fall back to NIC with default route
                $route = Get-NetRoute -DestinationPrefix "0.0.0.0/0" -ErrorAction SilentlyContinue |
                         Sort-Object RouteMetric | Select-Object -First 1
                if ($route) {
                    $adapter = Get-NetAdapter -InterfaceIndex $route.InterfaceIndex -ErrorAction SilentlyContinue
                    $targetNIC = $adapter.Name
                }
                if (-not $targetNIC) {
                    $targetNIC = (Get-NetAdapter -Physical | Where-Object Status -eq Up | Select-Object -First 1).Name
                }
            }
            Set-DnsClientServerAddress -InterfaceAlias $targetNIC -ServerAddresses $OldDNS -ErrorAction Stop
            Clear-DnsClientCache -ErrorAction SilentlyContinue
            $out.Status = "Success"
        }
        catch {
            $out.Status = "Failed"
            $out.Error  = $_.Exception.Message
        }
        return $out
    }

    # Attempt WinRM → PSSession → PsExec (same pattern as Set-WindowsDns)
    $remoteResult     = $null
    $connectionMethod = ""

    try {
        $sessionOpt   = New-PSSessionOption -OpenTimeout ($Config.WinRMTimeout * 1000)
        $invokeParams = @{
            ComputerName  = $hostname
            ScriptBlock   = $restoreBlock
            ArgumentList  = @(,$oldDNS), $nicName
            SessionOption = $sessionOpt
            ErrorAction   = "Stop"
        }
        if ($useCredential) { $invokeParams.Credential = $Credential }
        $remoteResult     = Invoke-Command @invokeParams
        $connectionMethod = "WinRM"
    }
    catch {
        Write-Warning "[$hostname] WinRM failed: $($_.Exception.Message)"
        try {
            $sessionOpt    = New-PSSessionOption -OpenTimeout ($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 $restoreBlock `
                                -ArgumentList @(,$oldDNS), $nicName
            Remove-PSSession $session -ErrorAction SilentlyContinue
            $connectionMethod = "PSSession"
        }
        catch {
            Write-Warning "[$hostname] PSSession failed: $($_.Exception.Message)"
            if (Test-Path $Config.PsExecPath) {
                $dnsJoined  = '"{0}"' -f ($oldDNS -join '","')
                $dnsArray   = "@($dnsJoined)"
                $nicArg     = if ($nicName) { "`$nic='$nicName'" } else {
                    "`$route=Get-NetRoute -DestinationPrefix '0.0.0.0/0'|Sort RouteMetric|Select -First 1;" +
                    "`$nic=(Get-NetAdapter -InterfaceIndex `$route.InterfaceIndex).Name"
                }
                $credArgs = @()
                if ($useCredential) {
                    $credArgs = @("-u", $Credential.UserName, "-p", $Credential.GetNetworkCredential().Password)
                }
                $cmd = "$nicArg; Set-DnsClientServerAddress -InterfaceAlias `$nic -ServerAddresses $dnsArray; Clear-DnsClientCache; Write-Output 'ROLLBACK_OK'"
                $enc = [Convert]::ToBase64String([Text.Encoding]::Unicode.GetBytes($cmd))
                $psexecArgs = @("\\$hostname", "-accepteula", "-nobanner", "-h") +
                              $credArgs +
                              @("powershell.exe", "-NonInteractive", "-EncodedCommand", $enc)
                $out = & $Config.PsExecPath @psexecArgs 2>&1

                $connectionMethod = "PsExec"
                $remoteResult = @{
                    Status = if ($out -match "ROLLBACK_OK") { "Success" } else { "Failed" }
                    Error  = if ($out -notmatch "ROLLBACK_OK") { $out -join " " } else { "" }
                }
            }
            else {
                $result.ChangeStatus     = "Failed"
                $result.ConnectionMethod = "None"
                $result.ErrorMessage     = "All connection methods failed and PsExec not found"
                return $result
            }
        }
    }

    $result.ConnectionMethod = $connectionMethod
    $result.ChangeStatus     = if ($remoteResult.Status -eq "Success") { "Success" } else { "Failed" }
    $result.ErrorMessage     = $remoteResult.Error

    # Validate — confirm old DNS IPs are resolvable
    if ($result.ChangeStatus -eq "Success") {
        $result.ValidationStatus = Test-WindowsDnsValidation `
            -Hostname         $hostname `
            -NewDNS           ($oldDNS | Select-Object -First 1) `
            -Domain           $Config.Domain `
            -ConnectionMethod $connectionMethod `
            -Config           $Config `
            -Credential       $Credential
    }

    return $result
}


# ── Helper: convert PSCustomObject rollback entry to hashtable ────────────────
function ConvertTo-RollbackHashtable {
    <#
    .SYNOPSIS
        Converts a JSON-deserialized rollback entry (PSCustomObject) to a
        hashtable with the shape expected by Invoke-LinuxDnsRollback.
    #>
    param([object]$Entry)

    # Normalise ModifiedFiles — JSON deserialization returns PSCustomObject[]
    $modifiedFiles = @($Entry.ModifiedFiles | ForEach-Object {
        @{
            Path          = $_.Path
            ContentBefore = $_.ContentBefore
            ContentAfter  = $_.ContentAfter
        }
    })

    return @{
        Hostname         = $Entry.Hostname
        OS               = $Entry.OS
        Distro           = $Entry.Distro
        ConnectionMethod = $Entry.ConnectionMethod
        OldDNS           = @($Entry.OldDNS)
        NewDNS           = @($Entry.NewDNS)
        NICName          = $Entry.NICName
        ModifiedFiles    = $modifiedFiles
        ServiceRestarted = $Entry.ServiceRestarted
        ResolvConfBefore = $Entry.ResolvConfBefore
        ResolvConfAfter  = $Entry.ResolvConfAfter
        Timestamp        = $Entry.Timestamp
    }
}