$script:WorkDir = ".\work"
$script:MountDir = ".\work\mount"
$script:MediaDir = ".\work\media"
$script:BootWimPath = ".\work\media\sources\boot.wim"

function Initialize-WinPEWorkspace {
    <#
    .SYNOPSIS
    Creates WinPE workspace using copype.cmd from Windows ADK.
    #>
    [CmdletBinding()]
    param()

    if (Test-Path $script:WorkDir) {
        Write-Host "Cleaning existing workspace..." -ForegroundColor Yellow
        Remove-Item $script:WorkDir -Recurse -Force
    }

    Write-Host "Initializing WinPE workspace..." -ForegroundColor Yellow

    $adkPath = "C:\Program Files (x86)\Windows Kits\10\Assessment and Deployment Kit"
    $copypePath = Join-Path $adkPath "Windows Preinstallation Environment\copype.cmd"
    $setEnvPath = Join-Path $adkPath "Deployment Tools\DandISetEnv.bat"

    if (-not (Test-Path $copypePath)) {
        throw "copype.cmd not found. Is Windows ADK installed?"
    }

    if (-not (Test-Path $setEnvPath)) {
        throw "DandISetEnv.bat not found. ADK installation may be incomplete."
    }

    $currentLocation = Get-Location
    $workPath = Join-Path $currentLocation "work"

    # Create a temporary batch file to run both commands in sequence
    $tempBatchFile = Join-Path $env:TEMP "winpe-init-$(Get-Random).cmd"

    try {
        # Write batch file that calls DandISetEnv, then copype
        $batchContent = @"
@echo off
call "$setEnvPath"
call "$copypePath" amd64 "$workPath"
exit /b %ERRORLEVEL%
"@
        Set-Content -Path $tempBatchFile -Value $batchContent -Encoding ASCII

        # Execute the batch file
        $process = Start-Process -FilePath "cmd.exe" `
            -ArgumentList "/c", $tempBatchFile `
            -Wait -PassThru -NoNewWindow

        $exitCode = $process.ExitCode
    }
    finally {
        # Clean up temp file
        if (Test-Path $tempBatchFile) {
            Remove-Item $tempBatchFile -Force -ErrorAction SilentlyContinue
        }
    }

    if ($exitCode -ne 0) {
        throw "Failed to initialize WinPE workspace. Exit code: $exitCode"
    }

    if (-not (Test-Path $script:BootWimPath)) {
        throw "boot.wim not found after workspace initialization"
    }

    Write-Host "  ✓ WinPE workspace initialized" -ForegroundColor Green
}

function Mount-WinPEImage {
    <#
    .SYNOPSIS
    Mounts the WinPE boot.wim image for modification.
    #>
    [CmdletBinding()]
    param()

    if (-not (Test-Path $script:MountDir)) {
        New-Item -ItemType Directory -Path $script:MountDir | Out-Null
    }

    Write-Host "Mounting WinPE image..." -ForegroundColor Yellow

    # Use DISM directly
    & dism.exe /Mount-Image /ImageFile:$script:BootWimPath /Index:1 /MountDir:$script:MountDir | Out-Host

    # Verify mount succeeded by checking for Windows directory in mount point
    $windowsDir = Join-Path $script:MountDir "Windows"
    if (-not (Test-Path $windowsDir)) {
        throw "Failed to mount WinPE image. Mount verification failed - Windows directory not found."
    }

    Write-Host "  ✓ WinPE image mounted at: $script:MountDir" -ForegroundColor Green
}

function Add-WinPEOptionalComponent {
    <#
    .SYNOPSIS
    Adds an optional component (and its language pack) to the mounted WinPE image.
    #>
    [CmdletBinding()]
    param(
        [Parameter(Mandatory = $true)]
        [string]$Name
    )

    $adkPath = "C:\Program Files (x86)\Windows Kits\10\Assessment and Deployment Kit"
    $ocPath = Join-Path $adkPath "Windows Preinstallation Environment\amd64\WinPE_OCs"

    $componentCab = Join-Path $ocPath "$Name.cab"
    $languageCab = Join-Path $ocPath "en-us\$Name`_en-us.cab"

    if (-not (Test-Path $componentCab)) {
        throw "Optional component not found: $componentCab"
    }

    Write-Host "  Adding $Name..." -ForegroundColor Cyan

    & dism.exe /Add-Package /Image:$script:MountDir /PackagePath:$componentCab
    if ($LASTEXITCODE -ne 0) {
        throw "Failed to add component $Name. Exit code: $LASTEXITCODE"
    }

    if (Test-Path $languageCab) {
        Write-Host "    Adding $Name language pack..." -ForegroundColor Cyan

        & dism.exe /Add-Package /Image:$script:MountDir /PackagePath:$languageCab
        if ($LASTEXITCODE -ne 0) {
            Write-Warning "Failed to add language pack for $Name (non-critical)"
        }
    }

    Write-Host "    ✓ $Name installed" -ForegroundColor Green
}

function Set-WinPEScratchSpace {
    <#
    .SYNOPSIS
    Sets the scratch space size for better WinPE performance.
    #>
    [CmdletBinding()]
    param(
        [Parameter(Mandatory = $false)]
        [int]$SizeMB = 512
    )

    Write-Host "Setting scratch space to ${SizeMB}MB..." -ForegroundColor Yellow

    & dism.exe /Set-ScratchSpace:$SizeMB /Image:$script:MountDir
    if ($LASTEXITCODE -ne 0) {
        Write-Warning "Failed to set scratch space (non-critical)"
    } else {
        Write-Host "  ✓ Scratch space configured" -ForegroundColor Green
    }
}

function Dismount-WinPEImage {
    <#
    .SYNOPSIS
    Unmounts and commits changes to the WinPE image.
    #>
    [CmdletBinding()]
    param()

    Write-Host "Unmounting WinPE image and committing changes..." -ForegroundColor Yellow

    & dism.exe /Unmount-Image /MountDir:$script:MountDir /Commit
    if ($LASTEXITCODE -ne 0) {
        Write-Error "Failed to unmount image. Attempting cleanup..."

        & dism.exe /Cleanup-Mountpoints

        throw "Failed to unmount WinPE image. Exit code: $LASTEXITCODE"
    }

    Write-Host "  ✓ WinPE image unmounted successfully" -ForegroundColor Green
}

function New-BootableMedia {
    <#
    .SYNOPSIS
    Creates bootable ISO file from WinPE workspace.
    #>
    [CmdletBinding()]
    param(
        [Parameter(Mandatory = $false)]
        [string]$OutputPath = ".\output"
    )

    if (-not (Test-Path $OutputPath)) {
        New-Item -ItemType Directory -Path $OutputPath | Out-Null
    }

    $isoPath = Join-Path $OutputPath "WinPE.iso"

    if (Test-Path $isoPath) {
        Write-Host "Removing existing ISO..." -ForegroundColor Yellow
        Remove-Item $isoPath -Force
    }

    Write-Host "Creating bootable ISO..." -ForegroundColor Yellow

    $adkPath = "C:\Program Files (x86)\Windows Kits\10\Assessment and Deployment Kit"
    $makeWinPEMediaPath = Join-Path $adkPath "Windows Preinstallation Environment\MakeWinPEMedia.cmd"
    $setEnvPath = Join-Path $adkPath "Deployment Tools\DandISetEnv.bat"

    if (-not (Test-Path $makeWinPEMediaPath)) {
        throw "MakeWinPEMedia.cmd not found"
    }

    if (-not (Test-Path $setEnvPath)) {
        throw "DandISetEnv.bat not found. ADK installation may be incomplete."
    }

    $currentLocation = Get-Location
    $workPath = Join-Path $currentLocation "work"

    # Create a temporary batch file - need to set up ADK environment first
    $tempBatchFile = Join-Path $env:TEMP "winpe-makeiso-$(Get-Random).cmd"

    try {
        # Write batch file that calls DandISetEnv, then MakeWinPEMedia
        $batchContent = @"
@echo off
call "$setEnvPath"
call "$makeWinPEMediaPath" /ISO "$workPath" "$isoPath"
exit /b %ERRORLEVEL%
"@
        Set-Content -Path $tempBatchFile -Value $batchContent -Encoding ASCII

        # Execute the batch file
        $process = Start-Process -FilePath "cmd.exe" `
            -ArgumentList "/c", $tempBatchFile `
            -Wait -PassThru -NoNewWindow

        $exitCode = $process.ExitCode
    }
    finally {
        # Clean up temp file
        if (Test-Path $tempBatchFile) {
            Remove-Item $tempBatchFile -Force -ErrorAction SilentlyContinue
        }
    }

    if ($exitCode -ne 0) {
        throw "Failed to create ISO. Exit code: $exitCode"
    }

    if (-not (Test-Path $isoPath)) {
        throw "ISO creation reported success but file not found"
    }

    $isoSize = (Get-Item $isoPath).Length / 1MB
    Write-Host "  ✓ Bootable ISO created: $isoPath ($([math]::Round($isoSize, 2)) MB)" -ForegroundColor Green
}

function Get-MountDirectory {
    <#
    .SYNOPSIS
    Returns the current mount directory path.
    #>
    return $script:MountDir
}

Export-ModuleMember -Function @(
    'Initialize-WinPEWorkspace',
    'Mount-WinPEImage',
    'Add-WinPEOptionalComponent',
    'Set-WinPEScratchSpace',
    'Dismount-WinPEImage',
    'New-BootableMedia',
    'Get-MountDirectory'
)