#Requires -RunAsAdministrator
param(
    [string]$IsoPath = "$PSScriptRoot\..\..\tiny10.iso",
    [string]$OutputIso = ""
)

$projectRoot = (Resolve-Path "$PSScriptRoot\..\..")

if (-not (Test-Path $IsoPath)) {
    Write-Error "ISO not found: $IsoPath"
    exit 1
}

$IsoPath = (Resolve-Path $IsoPath).Path
if (-not $OutputIso) { $OutputIso = $IsoPath }

# Get Windows ADK path from registry(following Visual Studio's winsdk.bat approach).
$WinSDKPath = [Microsoft.Win32.Registry]::GetValue("HKEY_LOCAL_MACHINE\SOFTWARE\WOW6432Node\Microsoft\Windows Kits\Installed Roots", "KitsRoot10", $null)
if ($null -eq $WinSDKPath) {
    $WinSDKPath = [Microsoft.Win32.Registry]::GetValue("HKEY_LOCAL_MACHINE\SOFTWARE\Microsoft\Windows Kits\Installed Roots", "KitsRoot10", $null)
}

if ($null -ne $WinSDKPath) {
    # Trim the following backslash for path concatenation.
    $WinSDKPath = $WinSDKPath.TrimEnd('\')
    $ADKDepTools = "$WinSDKPath\Assessment and Deployment Kit\Deployment Tools\$hostarchitecture\Oscdimg"
}
$localOSCDIMGPath = "$PSScriptRoot\oscdimg.exe"

if ((Test-Path variable:ADKDepTools) -and (Test-Path "$ADKDepTools\oscdimg.exe" -PathType leaf)) {
    Write-Host "Will be using oscdimg.exe from system ADK."
    $OSCDIMG = "$ADKDepTools\oscdimg.exe"
}
else {
    Write-Host "oscdimg.exe from system ADK not found. Will be using bundled oscdimg.exe."
    
    $url = "https://msdl.microsoft.com/download/symbols/oscdimg.exe/3D44737265000/oscdimg.exe"

    if (![System.IO.File]::Exists($localOSCDIMGPath)) {
        Write-Host "Downloading oscdimg.exe..."
        Invoke-WebRequest -Uri $url -OutFile $localOSCDIMGPath

        if ([System.IO.File]::Exists($localOSCDIMGPath)) {
            Write-Host "oscdimg.exe downloaded successfully."
        }
        else {
            Write-Error "Failed to download oscdimg.exe."
            exit 1
        }
    }
    else {
        Write-Host "oscdimg.exe already exists locally."
    }

    $OSCDIMG = $localOSCDIMGPath
}

# Temp working folder
$tempDir = "$env:TEMP\tiny10_patch_$([System.IO.Path]::GetRandomFileName().Replace('.',''))"
New-Item -ItemType Directory -Path $tempDir | Out-Null
Write-Host "Temp dir: $tempDir"

try {
    # Mount ISO
    Write-Host "Mounting ISO..."
    $diskImage = Mount-DiskImage -ImagePath $IsoPath -PassThru
    $driveLetter = ($diskImage | Get-Volume).DriveLetter + ":"
    Write-Host "Mounted at $driveLetter"

    # Copy contents
    Write-Host "Copying ISO contents (this may take a minute)..."
    Copy-Item -Path "$driveLetter\*" -Destination $tempDir -Recurse -Force

    Dismount-DiskImage -ImagePath $IsoPath | Out-Null
    Write-Host "Dismounted."

    # Clear read-only flags
    Get-ChildItem -Path $tempDir -Recurse | ForEach-Object {
        if ($_.Attributes -band [System.IO.FileAttributes]::ReadOnly) {
            $_.Attributes = $_.Attributes -bxor [System.IO.FileAttributes]::ReadOnly
        }
    }

    # Replace autounattend.xml
    $newXml = "$projectRoot\includes\autounattend-win10.xml"
    if (-not (Test-Path $newXml)) {
        Write-Error "autounattend-win10.xml not found at: $newXml"
        exit 1
    }
    Write-Host "Replacing autounattend.xml..."
    Copy-Item -Path $newXml -Destination "$tempDir\autounattend.xml" -Force

    # Detect label from ISO filename
    $isoName = [System.IO.Path]::GetFileNameWithoutExtension($IsoPath)
    $label = if ($isoName -match 'x86') { "Tiny10_x86" } else { "Tiny10_x64" }

    $bootEtfs = "$tempDir\boot\etfsboot.com"
    $bootEfi  = "$tempDir\efi\microsoft\boot\efisys.bin"
    Write-Host "Repacking ISO as: $OutputIso"
    & "$OSCDIMG" "-l$label" '-m' '-o' '-u2' '-udfver102' `
        "-bootdata:2#p0,e,b$bootEtfs#pEF,e,b$bootEfi" `
        $tempDir $OutputIso

    if ($LASTEXITCODE -eq 0) {
        Write-Host "Done. ISO saved to: $OutputIso"
    } else {
        Write-Error "oscdimg failed with exit code $LASTEXITCODE"
    }
} finally {
    try { Dismount-DiskImage -ImagePath $IsoPath -ErrorAction SilentlyContinue | Out-Null } catch {}
    Write-Host "Cleaning up temp dir..."
    Remove-Item -Path $tempDir -Recurse -Force -ErrorAction SilentlyContinue
}