# Auto-elevate to administrator
if (-not ([Security.Principal.WindowsPrincipal][Security.Principal.WindowsIdentity]::GetCurrent()).IsInRole([Security.Principal.WindowsBuiltInRole]::Administrator)) {
    $maxRetries = 5
    $retryCount = 0
    $launched = $false

    while ($retryCount -lt $maxRetries -and -not $launched) {
        $retryCount++
        try {
            Write-Host "[Elevate $retryCount/$maxRetries] Requesting admin privileges..." -ForegroundColor Cyan
            $proc = Start-Process powershell.exe "-NoProfile -ExecutionPolicy Bypass -File `"$PSCommandPath`"" -Verb RunAs -Wait -PassThru -ErrorAction Stop
            if ($null -ne $proc) {
                $launched = $true
            }
        }
        catch {
            Write-Host "Elevation failed: user cancelled or permission error" -ForegroundColor Yellow
            if ($retryCount -ge $maxRetries) {
                Write-Host "Max retries ($maxRetries) reached, exiting." -ForegroundColor Red
                exit 1
            }
            Start-Sleep -Milliseconds 300
        }
    }

    if ($launched) {
        exit 0
    }
    exit 1
}

Write-Host "Running as administrator!" -ForegroundColor Green

# ==========================================================
# ImDisk-based FAT32 disk image builder
# ==========================================================

$SCRIPT_DIR = Split-Path -Parent $MyInvocation.MyCommand.Definition
$PROJECT_DIR = Split-Path -Parent $SCRIPT_DIR
$DISK_IMG = Join-Path $PROJECT_DIR "disk.img"
$DISK_DIR = Join-Path $PROJECT_DIR "disk"
$MOUNT_LETTER = "Z:"

$IMAGE_SIZE = 256 * 1024 * 1024

Write-Host "=== Creating disk image ===" -ForegroundColor Cyan

if (-not (Get-Command "imdisk.exe" -ErrorAction SilentlyContinue)) {
    Write-Host "Error: Cannot find imdisk.exe. Please ensure ImDisk is installed and added to system PATH." -ForegroundColor Red
    exit 1
}

if (-not (Test-Path $DISK_DIR)) {
    Write-Host "Error: Source directory not found $DISK_DIR" -ForegroundColor Red
    exit 1
}

if (Test-Path $DISK_IMG) {
    Write-Host "Removing existing disk image..."
    Remove-Item -Force $DISK_IMG
}

Write-Host "Creating $IMAGE_SIZE byte disk image..."
$fsutilFile = [System.IO.File]::Create($DISK_IMG)
$fsutilFile.SetLength($IMAGE_SIZE)
$fsutilFile.Close()

if (Test-Path $MOUNT_LETTER) {
    Write-Host "Warning: $MOUNT_LETTER already exists. Attempting to detach..." -ForegroundColor Yellow
    imdisk -D -m $MOUNT_LETTER
    Start-Sleep -Seconds 1
}

Write-Host "Mounting and formatting disk image to $MOUNT_LETTER ..."
imdisk -a -f $DISK_IMG -m $MOUNT_LETTER -o hd --% -p "/fs:fat32 /q /y"

Start-Sleep -Seconds 2

Write-Host "Copying files to disk image..."
robocopy $DISK_DIR "$MOUNT_LETTER\" /E /NJH /NJS /np
if ($LASTEXITCODE -le 7) {
    Write-Host "Files copied successfully." -ForegroundColor Green
} else {
    Write-Host "Warning: File copy encountered an error. Robocopy exit code: $LASTEXITCODE" -ForegroundColor Yellow
}

Write-Host "Unmounting disk image..."
imdisk -D -m $MOUNT_LETTER

Write-Host "Disk image created: $DISK_IMG" -ForegroundColor Green

# ==========================================================
# Refresh the VMware disk next to the .vmx
# ==========================================================
# Vmware/ViperOS.vmdk is a monolithicFlat descriptor; the data lives in
# ViperOS-flat.vmdk.  Both are regenerated from disk.img so "Open VM" always
# sees the image that was just packed.  The .vmx already points at
# ide0:0.fileName = "ViperOS.vmdk"; it uses EFI firmware for bootx64.efi.
$VMWARE_DIR    = Join-Path $PROJECT_DIR "Vmware"
$VMWARE_VMDK   = Join-Path $VMWARE_DIR "ViperOS.vmdk"
$VMWARE_FLAT   = Join-Path $VMWARE_DIR "ViperOS-flat.vmdk"
$VMWARE_VMX    = Join-Path $VMWARE_DIR "Viper OS.vmx"

if (Test-Path $VMWARE_DIR) {
    Write-Host "=== Updating VMware disk ===" -ForegroundColor Cyan
    $vmRunning = Get-Process -Name "vmware-vmx" -ErrorAction SilentlyContinue
    if ($vmRunning) {
        Write-Host "Warning: a VMware VM is running; close it and re-run to refresh the VMDK." -ForegroundColor Yellow
    } else {
        # A leftover .lck from a crash keeps VMware from opening the VM.
        $vmLock = "$VMWARE_VMX.lck"
        if (Test-Path $vmLock) {
            Write-Host "Removing stale VMware lock: $vmLock" -ForegroundColor Yellow
            Remove-Item -Recurse -Force $vmLock -ErrorAction SilentlyContinue
        }
        $size    = (Get-Item $DISK_IMG).Length
        $sectors = [int64]($size / 512)
        $cyl     = [int64][math]::Floor($sectors / (16 * 63))

        # Keep the existing CID/UUID so VMware treats it as the same disk.
        $cid    = "fffffffe"
        $uuid   = "60 00 C2 96 88 c0 09 6d-c8 30 f7 71 57 ce 47 0d"
        $longid = "6af48e7def7fae6bfa57a4affffffffe"
        if (Test-Path $VMWARE_VMDK) {
            foreach ($l in (Get-Content $VMWARE_VMDK -ErrorAction SilentlyContinue)) {
                if ($l -match '^CID=([0-9a-fA-F]+)') { $cid = $Matches[1] }
                if ($l -match '^ddb\.uuid\s*=\s*"(.+)"') { $uuid = $Matches[1] }
                if ($l -match '^ddb\.longContentID\s*=\s*"(.+)"') { $longid = $Matches[1] }
            }
        }

        $descriptor = @"
# Disk DescriptorFile
version=1
encoding="GBK"
CID=$cid
parentCID=ffffffff
createType="monolithicFlat"

# Extent description
RW $sectors FLAT "ViperOS-flat.vmdk" 0

# The Disk Data Base
#DDB

ddb.adapterType = "ide"
ddb.geometry.cylinders = "$cyl"
ddb.geometry.heads = "16"
ddb.geometry.sectors = "63"
ddb.longContentID = "$longid"
ddb.uuid = "$uuid"
ddb.virtualHWVersion = "21"
"@

        try {
            if (Test-Path $VMWARE_FLAT) { Remove-Item -Force $VMWARE_FLAT }
            Copy-Item -Force $DISK_IMG $VMWARE_FLAT
            Set-Content -Path $VMWARE_VMDK -Value $descriptor -Encoding ASCII
            if (-not (Select-String -Path $VMWARE_VMX -Pattern 'ide0:0\.fileName\s*=\s*"ViperOS\.vmdk"' -Quiet)) {
                Write-Host "Note: $VMWARE_VMX does not reference ViperOS.vmdk; check ide0:0.fileName." -ForegroundColor Yellow
            }
            Write-Host "VMDK updated: $VMWARE_VMDK ($sectors sectors, $([math]::Round($size / 1MB)) MiB)" -ForegroundColor Green
        } catch {
            Write-Host "Warning: failed to refresh the VMware disk: $($_.Exception.Message)" -ForegroundColor Yellow
        }
    }
} else {
    Write-Host "Note: Vmware directory not found; skipping VMDK update." -ForegroundColor DarkGray
}