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

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

# Detect Windows 10 vs 11 from install.wim/install.esd so the right autounattend variant gets injected
function Get-WindowsVersionFromMedia {
    param([string]$MediaRoot)

    $srcBase = $MediaRoot
    if (-not (Test-Path "$srcBase\sources\install.wim") -and -not (Test-Path "$srcBase\sources\install.esd") -and (Test-Path "$srcBase\x64\sources")) {
        $srcBase = "$srcBase\x64"
    }

    $imagePath = if (Test-Path "$srcBase\sources\install.wim") { "$srcBase\sources\install.wim" }
                 elseif (Test-Path "$srcBase\sources\install.esd") { "$srcBase\sources\install.esd" }
                 else { $null }

    if (-not $imagePath) { return $null }

    try {
        $imgInfo = Get-WindowsImage -ImagePath $imagePath -Index 1
        $parts = $imgInfo.Version.Split('.')
        if ($parts.Count -ge 3) {
            $build = [int]$parts[2]
            if ($build -ge 22000) { return "11" }
            elseif ($build -ge 10240) { return "10" }
        }
    }
    catch {
        Write-Warning "Could not read Windows version from media: $($_.Exception.Message)"
    }
    return $null
}

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
        }
    }

    # Detect Windows version so we inject the matching autounattend variant
    Write-Host "Detecting Windows version from media..."
    $detectedVersion = Get-WindowsVersionFromMedia -MediaRoot $tempDir
    if (-not $detectedVersion) {
        Write-Host "Could not automatically detect Windows version from media."
        do {
            $detectedVersion = Read-Host "Enter Windows version (10 or 11)"
        } while ($detectedVersion -notin @("10", "11"))
    }
    else {
        Write-Host "Detected Windows $detectedVersion media."
    }

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

    # Detect label from ISO filename, tagged with the detected version
    $isoName = [System.IO.Path]::GetFileNameWithoutExtension($IsoPath)
    $label = if ($isoName -match 'x86') { "Tiny${detectedVersion}_x86" } else { "Tiny${detectedVersion}_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
}