#Requires -Version 5.1
<#
.SYNOPSIS
    HOST-side script: delivers and runs Setup-LinuxBuild2404.sh inside the Linux template VM.

.DESCRIPTION
    Automates provisioning of the Ubuntu 24.04 template VM from the host machine via SSH/SCP:
      1. Pre-flight validation: ssh/scp in PATH, SSH key present, VMX exists (if supplied),
         IPv4 address valid, TCP/22 reachable
      2. Auto power-on + guest IP detection via vmrun getGuestIPAddress (when -VMXPath provided)
      3. Add host key to known_hosts (idempotent, uses StrictHostKeyChecking=accept-new)
      4. Copy Setup-LinuxBuild2404.sh to /tmp on guest via scp + validate file size
      5. Execute Setup-LinuxBuild2404.sh via SSH (stdout/stderr streamed to host console)
      6. Post-setup remote validation batch check (user, sudo, tools, CI dirs, swap, sshd)
      7. Graceful shutdown + vmrun snapshot (when -TakeSnapshot or -VMXPath provided)
      Final: print summary

    Every validation step uses Assert-Step: prints [OK] on success, throws on failure.
    The snapshot is only taken after this script exits 0.

    PREREQUISITES (do these manually before running this script):
      a. The VM was created by Deploy-LinuxBuild2404.ps1 and cloud-init has completed
      b. The ci_build SSH public key ($SshKeyPath.pub) is already authorised in the guest
         (cloud-init in Deploy-LinuxBuild2404.ps1 injects the public key automatically)
      c. The VM NIC is on VMnet8 (NAT) with internet access

    AFTER THIS SCRIPT COMPLETES (exit 0, all validations passed):
      If -TakeSnapshot was NOT set and -VMXPath was NOT provided:
        1. Shut down the VM manually
        2. Take a snapshot named '$SnapshotName' in VMware Workstation
        3. Use -SnapshotName '$SnapshotName' when calling New-BuildVM.ps1

.PARAMETER VMXPath
    Full path to the template VM's .vmx file.
    When provided:
      - If the VM is powered off, the script starts it automatically via vmrun.
      - The guest IP is auto-detected via vmrun getGuestIPAddress (requires open-vm-tools).
      - At the end, the VM is shut down and a snapshot is taken automatically.
    Either -VMXPath or -VMIPAddress (or both) must be supplied.

.PARAMETER VMIPAddress
    IP address of the template VM on VMnet8 NAT.
    Optional when -VMXPath is provided (IP is auto-detected via open-vm-tools).
    Required when -VMXPath is omitted.
    Example: 192.168.79.131

.PARAMETER CloudUser
    SSH username in the guest VM. Default: ci_build.

.PARAMETER SshKeyPath
    Path to the SSH private key on the host. Default: F:\CI\keys\ci_linux.
    The corresponding public key is expected at $SshKeyPath.pub.

.PARAMETER SnapshotName
    Name of the snapshot to create. Default: BaseClean-Linux.
    Must match the snapshot name used in New-BuildVM.ps1.

.PARAMETER SkipUpdate
    Pass --skip-update to Setup-LinuxBuild2404.sh (skips apt-get upgrade).

.PARAMETER InstallDotNet
    Pass --dotnet to Setup-LinuxBuild2404.sh (installs .NET SDK).

.PARAMETER NoMingw
    Pass --no-mingw to Setup-LinuxBuild2404.sh (skips MinGW cross-compiler).
    MinGW is installed by default; use this only to skip it.

.PARAMETER TakeSnapshot
    Shut down the VM and take a vmrun snapshot after provisioning completes.
    Implied automatically when -VMXPath is provided (snapshot is always taken
    when the script manages the VM lifecycle via vmrun).

.PARAMETER VMwareWorkstationDir
    Installation directory of VMware Workstation on the host.
    Default: C:\Program Files (x86)\VMware\VMware Workstation

.EXAMPLE
    # Fully automated — VMX path only (power-on + IP detection + snapshot automatic):
    .\Prepare-LinuxBuild2404.ps1 -VMXPath 'F:\CI\Templates\LinuxBuild2404\LinuxBuild2404.vmx'

    # Explicit IP, skip update (VM already running, no snapshot):
    .\Prepare-LinuxBuild2404.ps1 -VMIPAddress '192.168.79.131' -SkipUpdate

    # VMX + .NET SDK + explicit snapshot:
    .\Prepare-LinuxBuild2404.ps1 `
        -VMXPath 'F:\CI\Templates\LinuxBuild2404\LinuxBuild2404.vmx' `
        -InstallDotNet -TakeSnapshot

.NOTES
    Requires:
      - OpenSSH client on the host (ssh.exe + scp.exe in PATH)
      - VMware Workstation with vmrun.exe (only needed when -VMXPath is supplied)
      - open-vm-tools installed in the guest (for vmrun getGuestIPAddress)
      - The SSH private key at $SshKeyPath with correct permissions
#>
[CmdletBinding()]
param(
    # Path to the VM .vmx — enables auto power-on + IP detection + auto snapshot.
    # Either -VMXPath or -VMIPAddress (or both) must be supplied.
    [string] $VMXPath = '',

    # Guest IP on VMnet8 NAT. Optional when -VMXPath provided (auto-detected via open-vm-tools).
    [string] $VMIPAddress = '',

    [string] $CloudUser = 'ci_build',

    [string] $SshKeyPath = 'F:\CI\keys\ci_linux',

    # Snapshot name (case-sensitive — must match the name used in New-BuildVM.ps1)
    [string] $SnapshotName = 'BaseClean-Linux',

    [switch] $SkipUpdate,

    [switch] $InstallDotNet,

    [switch] $NoMingw,

    # Shut down VM and take snapshot after provisioning (implied when -VMXPath is provided)
    [switch] $TakeSnapshot,

    [string] $VMwareWorkstationDir = 'C:\Program Files (x86)\VMware\VMware Workstation'
)

Set-StrictMode -Version Latest
$ErrorActionPreference = 'Stop'
$ProgressPreference   = 'SilentlyContinue'   # suppress Write-Progress (Test-NetConnection, IWR, etc.)

# ── Locate vmrun.exe ──────────────────────────────────────────────────────────
$vmrunExe = Join-Path $VMwareWorkstationDir 'vmrun.exe'
if (-not (Test-Path $vmrunExe)) {
    # Fallback: try the 64-bit install path or PATH
    $alt64 = 'C:\Program Files\VMware\VMware Workstation\vmrun.exe'
    if (Test-Path $alt64) {
        $vmrunExe = $alt64
    } else {
        $vmrunInPath = Get-Command vmrun -ErrorAction SilentlyContinue
        if ($vmrunInPath) {
            $vmrunExe = $vmrunInPath.Source
        } else {
            throw "[Prepare] vmrun.exe not found at '$vmrunExe'. Install VMware Workstation or add it to PATH."
        }
    }
}

# ── Validate input: need VMXPath or VMIPAddress ───────────────────────────────
if ($VMXPath -eq '' -and $VMIPAddress -eq '') {
    throw "[Prepare] Supply -VMXPath (auto power-on + IP detection) or -VMIPAddress (or both)."
}

# ── Auto power-on + IP detection (when VMXPath provided) ─────────────────────
if ($VMXPath -ne '') {
    if (-not (Test-Path $VMXPath)) { throw "[Prepare] VMX not found: '$VMXPath'" }

    $runningList = @(& $vmrunExe list 2>&1 |
        Where-Object { $_ -match '\.vmx$' } |
        ForEach-Object { $_.Trim() })

    if ($runningList -notcontains $VMXPath) {
        Write-Host "[Prepare] VM not running — starting: $VMXPath"
        & $vmrunExe start $VMXPath nogui
        if ($LASTEXITCODE -ne 0) { throw "[Prepare] vmrun start failed (exit $LASTEXITCODE)." }
        Write-Host "[Prepare] VM started. Waiting for guest OS to initialize..."
    } else {
        Write-Host "[Prepare] VM already running: $VMXPath"
    }

    if ($VMIPAddress -eq '') {
        Write-Host "[Prepare] Detecting guest IP via open-vm-tools (getGuestIPAddress -wait)..."
        $detectedIP = (& $vmrunExe getGuestIPAddress $VMXPath -wait 2>&1).Trim()
        if ($detectedIP -notmatch '^\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}$') {
            throw "[Prepare] Could not get guest IP. Got: '$detectedIP'. Ensure open-vm-tools is running in the guest."
        }
        $VMIPAddress = $detectedIP
        Write-Host "[Prepare] Guest IP: $VMIPAddress" -ForegroundColor Green
    }
}

# ── Helper functions ──────────────────────────────────────────────────────────

function Assert-Step {
    param(
        [Parameter(Mandatory)] [string]      $StepName,
        [Parameter(Mandatory)] [string]      $Description,
        [Parameter(Mandatory)] [scriptblock] $Test
    )
    $ok = $false
    try { $ok = [bool](& $Test) }
    catch { throw "[Prepare/$StepName] Validation threw on '$Description': $_" }
    if (-not $ok) { throw "[Prepare/$StepName] Validation failed: $Description" }
    Write-Host "  [OK] $Description" -ForegroundColor Green
}

function Invoke-SshCommand {
    param(
        [Parameter(Mandatory)] [string] $IP,
        [Parameter(Mandatory)] [string] $User,
        [Parameter(Mandatory)] [string] $KeyPath,
        [Parameter(Mandatory)] [string] $Command,
        [switch] $PassThru   # return stdout instead of printing
    )
    $sshArgs = @(
        '-i', $KeyPath,
        '-o', 'StrictHostKeyChecking=no',
        '-o', 'UserKnownHostsFile=NUL',
        '-o', 'LogLevel=ERROR',
        '-o', 'ConnectTimeout=10',
        '-o', 'BatchMode=yes',
        "$User@$IP",
        $Command
    )
    if ($PassThru) {
        $savedEap = $ErrorActionPreference; $ErrorActionPreference = 'Continue'
        $out = & ssh @sshArgs 2>&1
        $sshRc = $LASTEXITCODE
        $ErrorActionPreference = $savedEap
        if ($sshRc -ne 0) { throw "[SSH] Command failed (exit $sshRc): $Command`nOutput: $out" }
        return ($out | Where-Object { $_ -isnot [System.Management.Automation.ErrorRecord] })
    } else {
        $savedEap = $ErrorActionPreference; $ErrorActionPreference = 'Continue'
        & ssh @sshArgs
        $sshRc = $LASTEXITCODE
        $ErrorActionPreference = $savedEap
        if ($sshRc -ne 0) { throw "[SSH] Command failed (exit $sshRc): $Command" }
    }
}

# ── Step 1 — Pre-flight validation ───────────────────────────────────────────
Write-Host "[Prepare] Step 1 — Pre-flight validation"

Assert-Step 'PreFlight' 'ssh.exe found in PATH' {
    $null -ne (Get-Command ssh -ErrorAction SilentlyContinue)
}
Assert-Step 'PreFlight' 'scp.exe found in PATH' {
    $null -ne (Get-Command scp -ErrorAction SilentlyContinue)
}
Assert-Step 'PreFlight' "SSH private key exists: $SshKeyPath" {
    Test-Path $SshKeyPath -PathType Leaf
}
Assert-Step 'PreFlight' "SSH public key exists: $SshKeyPath.pub" {
    Test-Path "$SshKeyPath.pub" -PathType Leaf
}
if ($VMXPath -ne '') {
    Assert-Step 'PreFlight' "VMX file exists: $VMXPath" {
        Test-Path $VMXPath -PathType Leaf
    }
}

# ── Step 2 — Validate IP address ─────────────────────────────────────────────
Write-Host "[Prepare] Step 2 — Validate guest IP address"
Assert-Step 'IPValidation' "VMIPAddress '$VMIPAddress' matches IPv4 pattern" {
    $VMIPAddress -match '^\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}$'
}
Write-Host "  Guest IP: $VMIPAddress"

# ── Step 3 — Wait for SSH port 22 ────────────────────────────────────────────
Write-Host "[Prepare] Step 3 — Waiting for SSH port 22 on $VMIPAddress (timeout 180 s)..."
$sshDeadline = (Get-Date).AddSeconds(180)
$sshReady    = $false
while ((Get-Date) -lt $sshDeadline) {
    if (Test-NetConnection -ComputerName $VMIPAddress -Port 22 `
            -InformationLevel Quiet -WarningAction SilentlyContinue -ErrorAction SilentlyContinue) {
        $sshReady = $true
        break
    }
    Start-Sleep -Seconds 10
    Write-Host "  ... waiting for SSH..."
}
Assert-Step 'SSHPort' "$VMIPAddress TCP/22 reachable" { $sshReady }

# ── Step 4 — Add host key to known_hosts ─────────────────────────────────────
Write-Host "[Prepare] Step 4 — Adding host key to known_hosts (StrictHostKeyChecking=accept-new)..."
$savedEap = $ErrorActionPreference; $ErrorActionPreference = 'Continue'
& ssh -i $SshKeyPath -o StrictHostKeyChecking=no -o UserKnownHostsFile=NUL -o LogLevel=ERROR -o ConnectTimeout=10 `
    -o BatchMode=yes "$CloudUser@$VMIPAddress" 'echo connected' 2>&1 | Out-Null
$ErrorActionPreference = $savedEap
if ($LASTEXITCODE -ne 0) {
    throw "[Prepare Step 4] SSH connectivity test failed (exit $LASTEXITCODE). Check key, user, and IP."
}
Write-Host "  [OK] SSH connectivity confirmed, host key added."

# ── Step 5 — Copy Setup-LinuxBuild2404.sh into VM ────────────────────────────
Write-Host "[Prepare] Step 5 — Copying Setup-LinuxBuild2404.sh to guest..."

$scriptDir   = Split-Path -Parent $MyInvocation.MyCommand.Path
$setupScript = Join-Path $scriptDir 'Setup-LinuxBuild2404.sh'
$guestScript = '/tmp/Setup-LinuxBuild2404.sh'

if (-not (Test-Path $setupScript -PathType Leaf)) {
    throw "[Prepare Step 5] Setup-LinuxBuild2404.sh not found alongside this script: $setupScript"
}

# Strip UTF-8 BOM and normalize CRLF→LF before copying.
# PS 5.1 / VS Code on Windows may save .sh files with BOM or CRLF; both break bash.
$shBytes = [System.IO.File]::ReadAllBytes($setupScript)
if ($shBytes.Length -ge 3 -and $shBytes[0] -eq 0xEF -and $shBytes[1] -eq 0xBB -and $shBytes[2] -eq 0xBF) {
    Write-Host "  [WARN] UTF-8 BOM detected in Setup-LinuxBuild2404.sh — stripping before copy"
    $shBytes = $shBytes[3..($shBytes.Length - 1)]
}
$shText    = [System.Text.Encoding]::UTF8.GetString($shBytes) -replace "`r`n", "`n" -replace "`r", "`n"
$tmpScript = Join-Path $env:TEMP 'Setup-LinuxBuild2404-ci.sh'
[System.IO.File]::WriteAllText($tmpScript, $shText, [System.Text.UTF8Encoding]::new($false))
$setupScript = $tmpScript

Write-Host "  Source: $setupScript (BOM/CRLF-cleaned)"
Write-Host "  Dest:   $guestScript"

$savedEap = $ErrorActionPreference; $ErrorActionPreference = 'Continue'
& scp -i $SshKeyPath -o StrictHostKeyChecking=no -o UserKnownHostsFile=NUL -o LogLevel=ERROR `
    $setupScript "${CloudUser}@${VMIPAddress}:${guestScript}" 2>&1 | Out-Null
$scpRc = $LASTEXITCODE
$ErrorActionPreference = $savedEap
if ($scpRc -ne 0) { throw "[Prepare Step 5] scp failed (exit $scpRc)." }

# Validate: compare local size vs remote size (allow +-2 bytes for LF/CRLF variance)
$localSize  = (Get-Item $setupScript).Length
$remoteSizeRaw = (Invoke-SshCommand -IP $VMIPAddress -User $CloudUser -KeyPath $SshKeyPath `
    -Command "stat -c %s $guestScript" -PassThru).Trim()
$remoteSize = [int]$remoteSizeRaw

Assert-Step 'FileCopy' "guest script size matches local ($localSize bytes, remote $remoteSize bytes, tolerance 2)" {
    [Math]::Abs($remoteSize - $localSize) -le 2
}

# ── Step 6 — Execute setup script ────────────────────────────────────────────
Write-Host "[Prepare] Step 6 — Running Setup-LinuxBuild2404.sh in guest (may take several minutes)..."

$setupFlags = [System.Collections.Generic.List[string]]::new()
if ($SkipUpdate)    { $setupFlags.Add('--skip-update') }
if ($InstallDotNet) { $setupFlags.Add('--dotnet') }
if ($NoMingw)       { $setupFlags.Add('--no-mingw') }
$flagsStr  = $setupFlags -join ' '
$remoteCmd = "chmod +x $guestScript && sudo $guestScript $flagsStr".TrimEnd()

Write-Host "  Command: $remoteCmd"
Write-Host "------------------------------------------------------------"

& ssh -i $SshKeyPath -o StrictHostKeyChecking=no -o UserKnownHostsFile=NUL -o ConnectTimeout=5 `
    -o BatchMode=yes "$CloudUser@$VMIPAddress" $remoteCmd
$setupExitCode = $LASTEXITCODE

Write-Host "------------------------------------------------------------"
if ($setupExitCode -ne 0) {
    throw "[Prepare Step 6] Setup-LinuxBuild2404.sh exited $setupExitCode — see output above."
}
Write-Host "[Prepare] Setup script completed (exit 0)." -ForegroundColor Green

# ── Step 7 — Post-setup remote validation ────────────────────────────────────
Write-Host "[Prepare] Step 7 — Running post-setup validation..."

$validationScript = @'
set -e
( id ci_build | grep -q sudo || sudo -l -U ci_build 2>/dev/null | grep -q NOPASSWD ) && echo "UserSudo:OK" || echo "UserSudo:FAIL"
sudo -n true 2>/dev/null && echo "SudoNoPass:OK" || echo "SudoNoPass:FAIL"
for t in gcc g++ clang cmake python3 git 7z; do command -v "$t" > /dev/null && echo "Tool:$t:OK" || echo "Tool:$t:FAIL"; done
test -d /opt/ci/build && echo "Dir:build:OK" || echo "Dir:build:FAIL"
test -d /opt/ci/output && echo "Dir:output:OK" || echo "Dir:output:FAIL"
test -d /opt/ci/scripts && echo "Dir:scripts:OK" || echo "Dir:scripts:FAIL"
swapon --show | grep -q . && echo "Swap:ACTIVE:FAIL" || echo "Swap:off:OK"
systemctl is-active ssh > /dev/null && echo "SSHD:active:OK" || echo "SSHD:inactive:FAIL"
grep -q "^PasswordAuthentication no" /etc/ssh/sshd_config && echo "SSHPwdAuth:disabled:OK" || echo "SSHPwdAuth:enabled:FAIL"
systemctl is-enabled ci-report-ip.service > /dev/null 2>&1 && echo "CIReportIP:enabled:OK" || echo "CIReportIP:enabled:FAIL"
test -x /usr/local/bin/ci-report-ip.sh && echo "CIReportIPScript:exists:OK" || echo "CIReportIPScript:exists:FAIL"
'@

$results = Invoke-SshCommand -IP $VMIPAddress -User $CloudUser -KeyPath $SshKeyPath `
    -Command $validationScript -PassThru

Write-Host "[Prepare] Validation results:"
$results | ForEach-Object { Write-Host "  $_" }

$failures = @($results | Where-Object { $_ -match ':FAIL' })
if ($failures.Count -gt 0) {
    throw "[Prepare Step 7] Post-setup validation failed:`n$($failures -join "`n")"
}
Write-Host "[Prepare] All post-setup checks passed." -ForegroundColor Green

# ── Step 8 — Shutdown + snapshot ─────────────────────────────────────────────
if ($TakeSnapshot -or $VMXPath -ne '') {
    Write-Host "[Prepare] Step 8 — Sending graceful shutdown to guest..."
    try {
        Invoke-SshCommand -IP $VMIPAddress -User $CloudUser -KeyPath $SshKeyPath `
            -Command 'sudo shutdown -h now' -PassThru | Out-Null
    } catch {
        Write-Warning "[Prepare] Shutdown command may have failed (VM may have already started shutting down): $_"
    }

    if ($VMXPath -ne '') {
        Write-Host "[Prepare] Waiting for VM to power off (timeout 120 s)..."
        $deadline    = (Get-Date).AddSeconds(120)
        $poweredOff  = $false
        while ((Get-Date) -lt $deadline) {
            $listOut = & $vmrunExe list 2>&1 | Out-String
            if ($listOut -notmatch [regex]::Escape($VMXPath)) {
                $poweredOff = $true
                break
            }
            Start-Sleep -Seconds 5
            Write-Host "  ... waiting for shutdown..."
        }
        if (-not $poweredOff) { throw "[Prepare Step 8] VM did not power off within 120 s." }

        # Delete existing snapshot with the same name so vmrun can create a fresh one.
        $existingSnaps = (& $vmrunExe listSnapshots $VMXPath 2>&1 | Out-String)
        if ($existingSnaps -match [regex]::Escape($SnapshotName)) {
            Write-Host "[Prepare] Snapshot '$SnapshotName' already exists — deleting before recreating..."
            & $vmrunExe deleteSnapshot $VMXPath $SnapshotName
            if ($LASTEXITCODE -ne 0) { throw "[Prepare Step 8] vmrun deleteSnapshot failed (exit $LASTEXITCODE)." }
            Write-Host "[Prepare] Old snapshot deleted."
        }

        Write-Host "[Prepare] Taking snapshot '$SnapshotName'..."
        & $vmrunExe snapshot $VMXPath $SnapshotName
        if ($LASTEXITCODE -ne 0) { throw "[Prepare Step 8] vmrun snapshot failed (exit $LASTEXITCODE)." }

        Assert-Step 'Snapshot' "snapshot '$SnapshotName' listed in vmrun listSnapshots" {
            $snaps = (& $vmrunExe listSnapshots $VMXPath 2>&1 | Out-String)
            $snaps -match [regex]::Escape($SnapshotName)
        }
        Write-Host "[Prepare] Snapshot '$SnapshotName' created." -ForegroundColor Green
    } else {
        Write-Host "[Prepare] -VMXPath not provided — snapshot skipped. Shut down and snapshot manually."
    }
} else {
    Write-Host "[Prepare] -TakeSnapshot not set — leaving VM running. Shut down and snapshot manually when ready."
}

# ── Final — Print summary ─────────────────────────────────────────────────────
Write-Host ""
Write-Host "[Prepare] Summary" -ForegroundColor Cyan
Write-Host "  Guest IP:  $VMIPAddress"
Write-Host "  SSH key:   $SshKeyPath"
if ($VMXPath -ne '') { Write-Host "  VMX:       $VMXPath" }
if ($TakeSnapshot -or $VMXPath -ne '') {
    Write-Host "  Snapshot:  $SnapshotName — ready for New-BuildVM.ps1 -SnapshotName '$SnapshotName'"
} else {
    Write-Host "  Snapshot:  not taken — run manually then use -SnapshotName '$SnapshotName'"
}