#Requires -Version 5.1
#Requires -RunAsAdministrator
<#
.SYNOPSIS
    Builds a customized Windows PE image with PowerShell 7 and script launcher.

.DESCRIPTION
    Automates the complete process of creating a WinPE image that includes:
    - PowerShell 7 integration
    - Custom PowerShell modules
    - Automatic script launcher on boot
    - Detection of external media with .scripts folder

.PARAMETER SkipADKCheck
    Skip Windows ADK installation check (not recommended).

.PARAMETER CleanBuild
    Removes existing work directory for a clean build.

.PARAMETER DefenderOptimization
    Optimize Windows Defender for faster build performance:
    - None: No changes to Defender (default, safest)
    - DisableRealTime: Temporarily disable real-time protection (fastest, 50-70% improvement, restored after build)
    - AddExclusions: Add build directories to exclusions (balanced, 20-30% improvement, cleaned up after)

.EXAMPLE
    .\Build-PE.ps1
    Performs a standard build with all checks

.EXAMPLE
    .\Build-PE.ps1 -CleanBuild
    Removes previous build artifacts and starts fresh

.EXAMPLE
    .\Build-PE.ps1 -DefenderOptimization DisableRealTime
    Builds with real-time protection temporarily disabled for maximum speed

.EXAMPLE
    .\Build-PE.ps1 -DefenderOptimization AddExclusions
    Adds build directories to Defender exclusions during build

.NOTES
    Requires Windows 10/11 with administrator privileges
    First run will download and install Windows ADK (~4GB, 20-30 min)
#>

[CmdletBinding()]
param(
    [Parameter(Mandatory = $false)]
    [switch]$SkipADKCheck,

    [Parameter(Mandatory = $false)]
    [switch]$CleanBuild,

    [Parameter(Mandatory = $false)]
    [ValidateSet('None', 'DisableRealTime', 'AddExclusions')]
    [string]$DefenderOptimization = 'None'
)

$ErrorActionPreference = 'Stop'
$ProgressPreference = 'SilentlyContinue'

$Script:StartTime = Get-Date
$Script:LogPath = Join-Path $PSScriptRoot "output\build.log"

#region Logging

function Write-Log {
    param(
        [string]$Message,
        [string]$Level = "INFO"
    )

    $timestamp = Get-Date -Format "yyyy-MM-dd HH:mm:ss"
    $logMessage = "[$timestamp] [$Level] $Message"

    if (-not (Test-Path (Split-Path $Script:LogPath))) {
        New-Item -ItemType Directory -Path (Split-Path $Script:LogPath) -Force | Out-Null
    }

    Add-Content -Path $Script:LogPath -Value $logMessage
}

function Write-Status {
    param(
        [string]$Message,
        [string]$Type = "Info"
    )

    $color = switch ($Type) {
        "Success" { "Green" }
        "Warning" { "Yellow" }
        "Error" { "Red" }
        "Info" { "Cyan" }
        "Progress" { "Magenta" }
        default { "White" }
    }

    $prefix = switch ($Type) {
        "Success" { "[✓]" }
        "Warning" { "[!]" }
        "Error" { "[✗]" }
        "Info" { "[i]" }
        "Progress" { "[*]" }
        default { "[-]" }
    }

    Write-Host "$prefix " -ForegroundColor $color -NoNewline
    Write-Host $Message

    Write-Log -Message $Message -Level $Type.ToUpper()
}

function Write-Header {
    param([string]$Title)

    $line = "=" * 80
    Write-Host ""
    Write-Host $line -ForegroundColor Cyan
    Write-Host "  $Title" -ForegroundColor Yellow
    Write-Host $line -ForegroundColor Cyan
    Write-Host ""
}

#endregion

#region Initialization

Write-Host ""
Write-Host "================================================================================" -ForegroundColor Cyan
Write-Host "                     WinPE Script Launcher Builder" -ForegroundColor Yellow
Write-Host "                          Production Build System" -ForegroundColor White
Write-Host "================================================================================" -ForegroundColor Cyan
Write-Host ""

Write-Log "Build started" "INFO"
Write-Status "Build log: $Script:LogPath" "Info"
Write-Host ""

if (Test-Path $Script:LogPath) {
    Clear-Content $Script:LogPath
}

#endregion

#region Import Modules

Write-Header "Loading Build Modules"

$libModules = @(
    "Dependencies.psm1",
    "WinPEBuilder.psm1",
    "PowerShell7.psm1",
    "ModuleManager.psm1",
    "DefenderOptimization.psm1"
)

foreach ($module in $libModules) {
    $modulePath = Join-Path $PSScriptRoot "lib\$module"

    if (-not (Test-Path $modulePath)) {
        Write-Status "Module not found: $module" "Error"
        throw "Required module missing: $modulePath"
    }

    try {
        Import-Module $modulePath -Force -ErrorAction Stop
        Write-Status "Loaded: $module" "Success"
    }
    catch {
        Write-Status "Failed to load: $module" "Error"
        throw "Module import failed: $_"
    }
}

Write-Host ""

#endregion

#region Load Configuration

Write-Header "Loading Configuration"

$configPath = Join-Path $PSScriptRoot "config.json"

if (-not (Test-Path $configPath)) {
    Write-Status "Configuration file not found: $configPath" "Error"
    throw "config.json is required"
}

try {
    $config = Get-Content $configPath -Raw | ConvertFrom-Json
    Write-Status "Configuration loaded successfully" "Success"
    Write-Status "PowerShell 7 Version: $($config.powershell7Version)" "Info"
    Write-Status "Startup Script: $(if ($config.startupScript) {$config.startupScript} else {'Invoke-ScriptLauncher.ps1'})" "Info"
    Write-Status "Target Script (External): $($config.targetScript)" "Info"
    Write-Status "Fallback Script (Offline): $(if ($config.fallbackScript) {$config.fallbackScript} else {'Invoke-ScriptMenu.ps1'})" "Info"
    Write-Status "Scripts Folder: $($config.scriptsFolder)" "Info"
    Write-Status "Modules: $($config.powershellModules.Count) configured" "Info"
}
catch {
    Write-Status "Failed to parse configuration" "Error"
    throw "Configuration error: $_"
}

Write-Host ""

#endregion

#region Dependency Checks

Write-Header "Validating Dependencies"

Write-Status "Checking administrator privileges..." "Progress"
if (-not (Test-AdminPrivileges)) {
    Write-Status "Administrator privileges required" "Error"
    throw "Please run this script as Administrator"
}
Write-Status "Running with administrator privileges" "Success"

if (-not $SkipADKCheck) {
    Write-Status "Checking Windows ADK installation..." "Progress"

    if (-not (Test-WindowsADK)) {
        Write-Status "Windows ADK not found" "Warning"
        Write-Host ""
        Write-Host "The Windows Assessment and Deployment Kit (ADK) is required." -ForegroundColor Yellow
        Write-Host "This will download and install:" -ForegroundColor White
        Write-Host "  - Windows ADK (~150MB download)" -ForegroundColor Cyan
        Write-Host "  - WinPE Add-on (~5GB download)" -ForegroundColor Cyan
        Write-Host "  - Installation time: ~20-30 minutes" -ForegroundColor Cyan
        Write-Host ""

        $response = Read-Host "Install Windows ADK now? (Y/N)"

        if ($response -ne 'Y' -and $response -ne 'y') {
            Write-Status "Build cancelled by user" "Warning"
            exit 0
        }

        Install-WindowsADK
        Write-Status "Windows ADK installed successfully" "Success"
    }
    else {
        Write-Status "Windows ADK is installed" "Success"
    }
}

Write-Status "Checking PowerShell 7 archive..." "Progress"
$ps7FileName = "PowerShell-$($config.powershell7Version)-win-x64.zip"
$ps7Path = Join-Path $PSScriptRoot $ps7FileName

if (-not (Test-PowerShell7Archive -Config $config)) {
    Write-Status "PowerShell 7 archive not found" "Warning"
    Write-Status "Downloading PowerShell $($config.powershell7Version)..." "Progress"

    Get-PowerShell7Archive -Config $config

    Write-Status "PowerShell 7 downloaded" "Success"
}
else {
    Write-Status "PowerShell 7 archive found" "Success"
}

Write-Host ""

#endregion

#region Clean Build Option

if ($CleanBuild) {
    Write-Header "Clean Build"

    $workDir = Join-Path $PSScriptRoot "work"

    if (Test-Path $workDir) {
        Write-Status "Removing existing work directory..." "Progress"
        try {
            Remove-Item $workDir -Recurse -Force
            Write-Status "Work directory cleaned" "Success"
        }
        catch {
            Write-Status "Failed to clean work directory" "Warning"
            Write-Status "Error: $($_.Exception.Message)" "Warning"
        }
    }

    Write-Host ""
}

#endregion

#region Build Process

# Windows Defender Optimization
$defenderState = $null
$defenderExclusionsAdded = $false

if ($DefenderOptimization -ne 'None') {
    Write-Header "Windows Defender Optimization"

    if (-not (Test-DefenderAvailable)) {
        Write-Status "Windows Defender not available on this system" "Warning"
        Write-Status "Continuing without Defender optimization" "Info"
        Write-Host ""
    }
    else {
        try {
            if ($DefenderOptimization -eq 'DisableRealTime') {
                Write-Status "Mode: Temporarily disable real-time protection" "Info"
                Write-Status "Disabling Windows Defender real-time protection..." "Progress"

                $defenderState = Disable-DefenderRealTimeProtection

                if ($null -ne $defenderState) {
                    Write-Status "Real-time protection disabled for build duration" "Success"
                    Write-Status "Will be restored automatically when build completes" "Info"
                    Write-Status "Expected performance improvement: 50-70%" "Info"
                }
            }
            elseif ($DefenderOptimization -eq 'AddExclusions') {
                Write-Status "Mode: Add build directories to exclusions" "Info"
                Write-Status "Adding build directories to Defender exclusions..." "Progress"

                Add-DefenderExclusions -PSScriptRoot $PSScriptRoot
                $defenderExclusionsAdded = $true

                Write-Status "Exclusions added (will be removed after build)" "Success"
                Write-Status "Expected performance improvement: 20-30%" "Info"
            }
        }
        catch {
            Write-Status "Defender optimization failed: $($_.Exception.Message)" "Warning"
            Write-Status "Continuing without Defender optimization" "Info"
        }

        Write-Host ""
    }
}

try {
    Write-Header "Building WinPE Image"

    Write-Status "Step 1/10: Initializing WinPE workspace" "Progress"
    Initialize-WinPEWorkspace
    Write-Status "Workspace initialized" "Success"
    Write-Host ""

    Write-Status "Step 2/10: Mounting WinPE image" "Progress"
    Mount-WinPEImage
    $mountDir = Get-MountDirectory
    Write-Status "Image mounted at: $mountDir" "Success"
    Write-Host ""

    Write-Status "Step 3/10: Adding optional components" "Progress"
    $components = @(
        "WinPE-WMI",
        "WinPE-NetFX",
        "WinPE-Scripting",
        "WinPE-PowerShell",
        "WinPE-StorageWMI",
        "WinPE-DismCmdlets"
    )

    foreach ($component in $components) {
        Add-WinPEOptionalComponent -Name $component
    }
    Write-Status "Optional components installed" "Success"
    Write-Host ""

    Write-Status "Step 4/10: Setting scratch space" "Progress"
    Set-WinPEScratchSpace -SizeMB 512
    Write-Host ""

    Write-Status "Step 5/10: Installing PowerShell 7" "Progress"
    Install-PowerShell7ToWinPE -ZipPath $ps7Path -MountDir $mountDir
    Write-Host ""

    Write-Status "Step 6/10: Copying PowerShell modules" "Progress"
    if ($config.powershellModules -and $config.powershellModules.Count -gt 0) {
        foreach ($moduleName in $config.powershellModules) {
            $modulePath = Get-RepositoryModulePath -ModuleName $moduleName

            Install-ModuleToWinPE -ModuleName $moduleName -SourcePath $modulePath -MountDir $mountDir
        }
        Write-Status "All modules copied successfully" "Success"
    }
    else {
        Write-Status "No modules configured - skipping" "Info"
    }
    Write-Host ""

    Write-Status "Step 7/10: Copying launcher scripts and configuration" "Progress"

    # Copy startup script
    $startupScriptName = if ($config.startupScript) { $config.startupScript } else { "Invoke-ScriptLauncher.ps1" }
    $startupSource = Join-Path $PSScriptRoot "resources\$startupScriptName"
    $startupDest = Join-Path $mountDir "Windows\System32\$startupScriptName"

    if (Test-Path $startupSource) {
        Copy-Item -Path $startupSource -Destination $startupDest -Force
        Write-Status "Startup script copied: $startupScriptName" "Success"
    }
    else {
        Write-Status "Startup script not found: $startupScriptName" "Error"
        throw "Startup script missing: $startupSource"
    }

    # Copy fallback script
    $fallbackScriptName = if ($config.fallbackScript) { $config.fallbackScript } else { "Invoke-ScriptMenu.ps1" }
    $fallbackSource = Join-Path $PSScriptRoot "resources\$fallbackScriptName"
    $fallbackDest = Join-Path $mountDir "Windows\System32\$fallbackScriptName"

    if (Test-Path $fallbackSource) {
        Copy-Item -Path $fallbackSource -Destination $fallbackDest -Force
        Write-Status "Fallback script copied: $fallbackScriptName" "Success"
    }
    else {
        Write-Status "Fallback script not found: $fallbackScriptName" "Error"
        throw "Fallback script missing: $fallbackSource"
    }

    # Copy configuration file
    $configDest = Join-Path $mountDir "Windows\System32\winpe-config.json"
    Copy-Item -Path $configPath -Destination $configDest -Force
    Write-Status "Configuration copied to image" "Success"
    Write-Host ""

    Write-Status "Step 8/10: Configuring startup script" "Progress"
    Set-PowerShell7Environment -MountDir $mountDir -LauncherScriptPath $startupScriptName
    Write-Host ""

    Write-Status "Step 9/10: Unmounting and committing changes" "Progress"
    Dismount-WinPEImage
    Write-Host ""

    Write-Status "Step 10/10: Creating bootable ISO" "Progress"
    New-BootableMedia -OutputPath (Join-Path $PSScriptRoot "output")

    Write-Host ""
    Write-Host "================================================================================" -ForegroundColor Green
    Write-Status "BUILD COMPLETED SUCCESSFULLY!" "Success"
    Write-Host "================================================================================" -ForegroundColor Green
    Write-Host ""

    $outputISO = Join-Path $PSScriptRoot "output\WinPE.iso"
    $isoSize = (Get-Item $outputISO).Length / 1MB

    Write-Status "Output: $outputISO" "Success"
    Write-Status "Size: $([math]::Round($isoSize, 2)) MB" "Info"

    $elapsed = (Get-Date) - $Script:StartTime
    Write-Status "Build time: $($elapsed.ToString('mm\:ss'))" "Info"
    Write-Host ""

    Write-Host "Next steps:" -ForegroundColor Yellow
    Write-Host "  1. Test ISO in virtual machine (Hyper-V, VirtualBox, etc.)" -ForegroundColor Cyan
    Write-Host "  2. Create .scripts folder on external media" -ForegroundColor Cyan
    Write-Host "  3. Add your scripts to .scripts folder" -ForegroundColor Cyan
    Write-Host "  4. Boot the WinPE image" -ForegroundColor Cyan
    Write-Host ""

    Write-Log "Build completed successfully" "INFO"
}
catch {
    Write-Host ""
    Write-Host "================================================================================" -ForegroundColor Red
    Write-Status "BUILD FAILED!" "Error"
    Write-Host "================================================================================" -ForegroundColor Red
    Write-Host ""
    Write-Status "Error: $($_.Exception.Message)" "Error"
    Write-Host ""
    Write-Status "Check log file for details: $Script:LogPath" "Info"
    Write-Host ""

    Write-Log "Build failed: $($_.Exception.Message)" "ERROR"

    try {
        Write-Status "Attempting to unmount image..." "Warning"
        Dismount-WinPEImage
    }
    catch {
        Write-Status "Could not unmount image cleanly" "Warning"
        Write-Status "Run 'dism /Cleanup-Mountpoints' to clean up" "Info"
    }

    exit 1
}
finally {
    # Restore Windows Defender settings
    if ($DefenderOptimization -ne 'None') {
        Write-Host ""
        Write-Header "Restoring Windows Defender Settings"

        try {
            if ($DefenderOptimization -eq 'DisableRealTime' -and $null -ne $defenderState) {
                Write-Status "Restoring Windows Defender real-time protection..." "Progress"
                Restore-DefenderRealTimeProtection -OriginalState $defenderState
                Write-Status "Real-time protection restored" "Success"
            }
            elseif ($DefenderOptimization -eq 'AddExclusions' -and $defenderExclusionsAdded) {
                Write-Status "Removing temporary Defender exclusions..." "Progress"
                Remove-DefenderExclusions -PSScriptRoot $PSScriptRoot
                Write-Status "Exclusions removed" "Success"
            }
        }
        catch {
            Write-Status "Failed to restore Defender settings: $($_.Exception.Message)" "Warning"
            Write-Status "You may need to manually restore Defender settings" "Warning"
        }

        Write-Host ""
    }
}

#endregion