param(
    [Nullable[int]]$Count,
    [string]$Region
)

try {
    Set-ExecutionPolicy -Scope Process -ExecutionPolicy Bypass -Force -ErrorAction Stop
} catch {
    # Ignore failures (e.g. already bypassed)
}

$ErrorActionPreference = 'Stop'

function Build-ForwardArgs {
    $argsList = @()

    if ($PSBoundParameters.ContainsKey('Count')) {
        $argsList += '-Count'
        $argsList += $Count
    }

    if ($PSBoundParameters.ContainsKey('Region')) {
        $argsList += '-Region'
        $argsList += $Region
    }

    if ($args.Count -gt 0) {
        $argsList += $args
    }

    if ($MyInvocation.UnboundArguments -and $MyInvocation.UnboundArguments.Count -gt 0) {
        $argsList += $MyInvocation.UnboundArguments
    }

    return $argsList
}

function Test-IsAdmin {
    return ([Security.Principal.WindowsPrincipal] `
        [Security.Principal.WindowsIdentity]::GetCurrent()
    ).IsInRole([Security.Principal.WindowsBuiltInRole]::Administrator)
}

if (-not (Test-IsAdmin)) {
    Write-Host "Re-launching with administrative privileges..." -ForegroundColor Yellow

    $forwardArgs = Build-ForwardArgs

    if ($PSCommandPath) {
        $startArgs = @("-NoProfile", "-ExecutionPolicy", "Bypass", "-File", $PSCommandPath) + $forwardArgs
        Start-Process powershell.exe -ArgumentList $startArgs -Verb RunAs
    }
    else {
        $exePath = (Get-Process -Id $PID).Path
        if ($forwardArgs.Count -gt 0) {
            Start-Process -FilePath $exePath -ArgumentList $forwardArgs -Verb RunAs
        }
        else {
            Start-Process -FilePath $exePath -Verb RunAs
        }
    }

    exit
}

# -----------------------------
# Config
# -----------------------------

$MinVersion = [Version]'7.6.0'
$ScriptDir = if ($PSScriptRoot -and (Test-Path $PSScriptRoot)) {
    $PSScriptRoot
}
elseif ($MyInvocation.MyCommand.Path) {
    Split-Path -Path $MyInvocation.MyCommand.Path -Parent
}
else {
    $PWD.Path
}

$PayloadScript = Join-Path $ScriptDir 'Payload.ps1'
$BundledMsi = Join-Path $ScriptDir 'PowerShell-7.6.3-win-x64.msi'

# Prefer the smallest Graph module required by the payload.
# The full Microsoft.Graph module is a meta-module and can be large, slow, and prone to lock/in-use issues.
$RequiredModules = @(
    'Microsoft.Graph.Authentication',`
    'Microsoft.Graph.Users',`
    'Microsoft.Graph.Identity.DirectoryManagement',`
    'Microsoft.Graph.Users.Actions',`
    'Microsoft.Graph.Identity.SignIns'
)

# If your payload needs additional Graph APIs, add only the specific submodules it requires.
# Example: @('Microsoft.Graph.Authentication', 'Microsoft.Graph.Users')

# -----------------------------
# Logging
# -----------------------------

$LogDir = "$env:ProgramData\LabBootstrap"
$LogFile = "$LogDir\bootstrap.log"

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

if (Test-Path $LogFile) {
    Remove-Item $LogFile -Force
}

function Log {
    param($msg)
    $line = "$(Get-Date -Format s) :: $msg"
    $line | Tee-Object -FilePath $LogFile -Append
}

function Log-PassedParameters {
    $scriptParams = Get-Variable -Name PSBoundParameters -Scope 1 -ValueOnly
    $scriptArgs = Get-Variable -Name args -Scope 1 -ValueOnly
    $scriptInvocation = Get-Variable -Name MyInvocation -Scope 1 -ValueOnly

    $boundParams = if ($scriptParams.Count -gt 0) {
        $scriptParams.GetEnumerator() | ForEach-Object { "$($_.Key)=$($_.Value)" } | Sort-Object
    } else {
        '<none>'
    }

    $argList = if ($scriptArgs.Count -gt 0) { $scriptArgs -join ' ' } else { '<none>' }
    $unboundArgs = if ($scriptInvocation.UnboundArguments -and $scriptInvocation.UnboundArguments.Count -gt 0) {
        $scriptInvocation.UnboundArguments -join ' '
    } else {
        '<none>'
    }

    Log "Passed named parameters: $($boundParams -join '; ')"
    Log "Passed positional args: $argList"
    Log "Passed unbound invocation args: $unboundArgs"
    Log "Raw command line: $([Environment]::CommandLine)"
    Log "Raw command line args: $([Environment]::GetCommandLineArgs() -join ' ')"
}

Log "==== Bootstrap started ===="
Log "Running as admin: $(Test-IsAdmin)"
Log-PassedParameters

# -----------------------------
# PowerShell detection
# -----------------------------

function Get-Pwsh {
    $results = @()
    $cmd = Get-Command pwsh -ErrorAction SilentlyContinue
    if ($cmd) { $results += $cmd.Source }

    $programFiles = @($env:ProgramW6432, $env:ProgramFiles) | Where-Object { $_ }
    foreach ($base in $programFiles) {
        $results += Get-ChildItem "$base\PowerShell\*\pwsh.exe" -ErrorAction SilentlyContinue |
            Sort-Object DirectoryName -Descending |
            Select-Object -ExpandProperty FullName
    }

    return $results | Select-Object -First 1
}

function Test-PwshValid {
    param($path)

    if (-not $path) { return $false }

    try {
        $v = & $path -NoProfile -Command '$PSVersionTable.PSVersion.ToString()'
        return ([Version]$v -ge $MinVersion)
    }
    catch {
        return $false
    }
}

function Install-PowerShell {
    param($msiPath)

    if (Test-Path $msiPath) {
        Log "Installing PowerShell from MSI: $msiPath"
        $proc = Start-Process msiexec.exe `
            -ArgumentList "/i", "$msiPath", "/qn", "/norestart", "ADD_PATH=1" `
            -Wait -PassThru
        if ($proc.ExitCode -ne 0) {
            Log "msiexec failed with exit code $($proc.ExitCode)"
            return $false
        }
        return $true
    }

    if (Get-Command winget -ErrorAction SilentlyContinue) {
        Log "Installing PowerShell via winget"
        $proc = Start-Process winget `
            -ArgumentList "install", "--id", "Microsoft.PowerShell", "--source", "winget", "--silent", "--accept-package-agreements", "--accept-source-agreements" `
            -Wait -PassThru
        if ($proc.ExitCode -ne 0) {
            Log "winget failed with exit code $($proc.ExitCode)"
            return $false
        }
        return $true
    }

    Log "No installer available: MSI not found and winget not installed"
    return $false
}

# -----------------------------
# Ensure PS7.6+
# -----------------------------

$pwsh = Get-Pwsh
if (-not (Test-PwshValid $pwsh)) {
    Log "PowerShell 7.6+ not found, installing..."
    if (-not (Install-PowerShell $BundledMsi)) {
        Log "ERROR: PowerShell installation failed"
        exit 1
    }

    Start-Sleep -Seconds 2
    $pwsh = Get-Pwsh
    if (-not (Test-PwshValid $pwsh)) {
        Log "ERROR: PowerShell installation failed after install"
        exit 1
    }
}

Log "Using pwsh: $pwsh"

# -----------------------------
# Install Modules (runs in pwsh)
# -----------------------------

$moduleScript = @'
$ErrorActionPreference = 'Stop'

if (-not (Get-PSRepository -Name PSGallery -ErrorAction SilentlyContinue)) {
    try {
        Register-PSRepository -Default -ErrorAction Stop
    } catch {
        # Ignore if repository is already registered or if the operation fails
    }
}

try {
    Set-PSRepository -Name PSGallery -InstallationPolicy Trusted -ErrorAction Stop
} catch {
    # Continue even if trust setting fails
}

foreach ($m in MODULE_LIST) {
    Write-Output "STEP: Checking for module $m"
    if (-not (Get-Module -ListAvailable -Name $m)) {
        Write-Output "STEP: Installing NuGet package provider"
        if (-not (Get-PackageProvider -Name NuGet -ErrorAction SilentlyContinue)) {
            Install-PackageProvider -Name NuGet `
                -ForceBootstrap `
                -MinimumVersion 2.8.5.201 `
                -Scope CurrentUser `
                -Force `
                -Confirm:$false `
                -ErrorAction Stop
        }

        Write-Output "STEP: Installing module: $m"
        Install-Module -Name $m `
            -Scope CurrentUser `
            -Repository PSGallery `
            -Force `
            -AllowClobber `
            -AcceptLicense `
            -Confirm:$false `
            -ErrorAction Stop
    }

    Write-Output "STEP: Verifying installed module $m"
    Get-Module -ListAvailable -Name $m | ForEach-Object {
        Write-Output "Found module: $($_.Name) $($_.Version) at $($_.Path)"
    }
}
'@

$moduleScript = $moduleScript -replace 'MODULE_LIST', "@('$($RequiredModules -join "','")')"

$tempScript = Join-Path $env:TEMP 'install-modules.ps1'
$moduleScript | Set-Content $tempScript -Encoding UTF8

Log "Starting module install script via pwsh: $tempScript"
$moduleOutput = & $pwsh `
    -NoProfile `
    -NoLogo `
    -NonInteractive `
    -ExecutionPolicy Bypass `
    -File $tempScript 2>&1
$moduleOutput | Tee-Object -FilePath $LogFile -Append
Log "Module install script finished"
if ($LASTEXITCODE -ne 0) {
    Log "Module installation failed with exit code $LASTEXITCODE"
    exit 1
}

# -----------------------------
# Run payload
# -----------------------------

if (-not (Test-Path $PayloadScript)) {
    $PayloadScript = Join-Path $env:TEMP 'Payload.ps1'
}

if (-not (Test-Path $PayloadScript)) {
    Log "ERROR: Payload not found: $PayloadScript"
    exit 1
}

$payloadParamEntries = @()
if ($Count -ne $null) {
    $payloadParamEntries += "Count = $Count"
}

if ($Region) {
    $payloadParamEntries += "Region = '$($Region -replace "'", "''")'"
}

$payloadSuffix = @()
if ($Count -ne $null) { $payloadSuffix += "-Count $Count" }
if ($Region) { $payloadSuffix += "-Region $Region" }
$payloadSuffix = if ($payloadSuffix.Count -gt 0) { ' with ' + ($payloadSuffix -join ' ') } else { '' }

Log "Running payload: $PayloadScript$payloadSuffix"
$payloadParamsLiteral = if ($payloadParamEntries.Count -gt 0) {
    "@{`r`n    " + ($payloadParamEntries -join "`r`n    ") + "`r`n}"
} else {
    '@{}'
}

$payloadScriptLiteral = "'" + ($PayloadScript -replace "'", "''") + "'"
$logFileLiteral = "'" + ($LogFile -replace "'", "''") + "'"
$wrapperScript = Join-Path $env:TEMP "run-payload-$([guid]::NewGuid()).ps1"
$wrapperContent = @"
`$payloadScript = $payloadScriptLiteral
`$payloadParams = $payloadParamsLiteral
`$transcriptPath = $logFileLiteral
`$exitCode = 0

try {
    Start-Transcript -Path `$transcriptPath -Append -Force -ErrorAction Stop | Out-Null
    Write-Host "Payload transcript started: `$transcriptPath" -ForegroundColor DarkGray
    & `$payloadScript @payloadParams
    if (-not `$?) {
        `$exitCode = 1
    }
}
catch {
    Write-Error "Payload wrapper caught error: `$(`$_.Exception.Message)"
    `$exitCode = 1
}
finally {
    try {
        Stop-Transcript | Out-Null
    }
    catch {}
}

exit `$exitCode
"@

$wrapperContent | Set-Content -LiteralPath $wrapperScript -Encoding UTF8
$payloadArgsList = @("-NoProfile", "-ExecutionPolicy", "Bypass", "-File", $wrapperScript)

$payloadProc = Start-Process $pwsh `
    -ArgumentList $payloadArgsList `
    -PassThru `
    -Wait `
    -NoNewWindow:$false

Remove-Item -LiteralPath $wrapperScript -Force -ErrorAction SilentlyContinue

if ($payloadProc.ExitCode -ne 0) {
    Log "Payload failed with exit code $($payloadProc.ExitCode)"
    exit 1
}

Log "==== Bootstrap complete ===="
