[CmdletBinding()]
param(
    [switch] $LoadFunctionsOnly,
    [string] $ExpectedUserSid
)

$ErrorActionPreference = "Stop"

$baseUrl = "https://julianbaumgardt.com"
$bootstrapUrl = "$baseUrl/w11"
$scriptUrl = "$baseUrl/w11-optimiser/w11-optimiser.ps1"
$manifestUrl = "$baseUrl/w11-optimiser/w11-optimiser.manifest.json"
$script:InstallRoot = $null
$script:ScriptPath = $null
$script:ManifestPath = $null
$script:ExpectedScriptHash = $null
$script:W11LastExitCode = 0
$script:W11Version = "unknown"

function Test-IsAdministrator {
    $identity = [Security.Principal.WindowsIdentity]::GetCurrent()
    $principal = [Security.Principal.WindowsPrincipal] $identity
    return $principal.IsInRole([Security.Principal.WindowsBuiltInRole]::Administrator)
}

function Get-CurrentUserSid {
    return [Security.Principal.WindowsIdentity]::GetCurrent().User.Value
}

function Assert-ExpectedUserIdentity {
    param([string] $ExpectedSid)

    if ([string]::IsNullOrWhiteSpace($ExpectedSid)) {
        return
    }
    try {
        $validatedSid = (New-Object -TypeName Security.Principal.SecurityIdentifier -ArgumentList $ExpectedSid).Value
    }
    catch {
        throw "Refused invalid initiating-user SID."
    }
    if ((Get-CurrentUserSid) -ne $validatedSid) {
        throw "Elevation changed Windows accounts. Re-run from a non-elevated PowerShell and approve UAC with the same account."
    }
}

function Assert-HttpsReleaseUri {
    param([Parameter(Mandatory = $true)][string] $Uri)

    $parsed = [Uri] $Uri
    if (-not $parsed.IsAbsoluteUri -or $parsed.Scheme -ne "https" -or
        $parsed.Host -notin @("julianbaumgardt.com", "www.julianbaumgardt.com")) {
        throw "Refused non-HTTPS or unapproved W11 Optimiser release URI: $Uri"
    }
}

function Set-AdminOnlyDirectoryAcl {
    param([Parameter(Mandatory = $true)][string] $Path)

    $inheritance = [Security.AccessControl.InheritanceFlags]::ContainerInherit -bor
        [Security.AccessControl.InheritanceFlags]::ObjectInherit
    $security = New-Object -TypeName Security.AccessControl.DirectorySecurity
    $security.SetAccessRuleProtection($true, $false)
    $administratorsSid = New-Object -TypeName Security.Principal.SecurityIdentifier -ArgumentList "S-1-5-32-544"
    $security.SetOwner($administratorsSid)
    foreach ($sidText in @("S-1-5-18", "S-1-5-32-544")) {
        $sid = New-Object -TypeName Security.Principal.SecurityIdentifier -ArgumentList $sidText
        $rule = New-Object -TypeName Security.AccessControl.FileSystemAccessRule -ArgumentList @(
            $sid,
            [Security.AccessControl.FileSystemRights]::FullControl,
            $inheritance,
            [Security.AccessControl.PropagationFlags]::None,
            [Security.AccessControl.AccessControlType]::Allow
        )
        [void] $security.AddAccessRule($rule)
    }
    [IO.Directory]::SetAccessControl($Path, $security)
}

function Set-AdminOnlyFileAcl {
    param([Parameter(Mandatory = $true)][string] $Path)

    $security = New-Object -TypeName Security.AccessControl.FileSecurity
    $security.SetAccessRuleProtection($true, $false)
    $administratorsSid = New-Object -TypeName Security.Principal.SecurityIdentifier -ArgumentList "S-1-5-32-544"
    $security.SetOwner($administratorsSid)
    foreach ($sidText in @("S-1-5-18", "S-1-5-32-544")) {
        $sid = New-Object -TypeName Security.Principal.SecurityIdentifier -ArgumentList $sidText
        $rule = New-Object -TypeName Security.AccessControl.FileSystemAccessRule -ArgumentList @(
            $sid,
            [Security.AccessControl.FileSystemRights]::FullControl,
            [Security.AccessControl.AccessControlType]::Allow
        )
        [void] $security.AddAccessRule($rule)
    }
    [IO.File]::SetAccessControl($Path, $security)
}

function Assert-AdminOnlyPathAcl {
    param([Parameter(Mandatory = $true)][string] $Path)

    $acl = Get-Acl -LiteralPath $Path -ErrorAction Stop
    $administratorsSid = "S-1-5-32-544"
    $allowedSids = @("S-1-5-18", $administratorsSid)
    if (-not $acl.AreAccessRulesProtected -or
        $acl.GetOwner([Security.Principal.SecurityIdentifier]).Value -ne $administratorsSid) {
        throw "Refused staging path without protected Administrator ownership: $Path"
    }
    $present = @{}
    foreach ($rule in $acl.GetAccessRules($true, $true, [Security.Principal.SecurityIdentifier])) {
        $sid = $rule.IdentityReference.Value
        if ($rule.AccessControlType -ne [Security.AccessControl.AccessControlType]::Allow -or $sid -notin $allowedSids) {
            throw "Refused staging path with an unexpected access rule: $Path"
        }
        $present[$sid] = $true
    }
    if (-not $present.ContainsKey("S-1-5-18") -or -not $present.ContainsKey($administratorsSid)) {
        throw "Refused staging path without SYSTEM and Administrators access: $Path"
    }
}

function Assert-DirectoryIsNotReparsePoint {
    param([Parameter(Mandatory = $true)][string] $Path)

    if ((Get-Item -LiteralPath $Path -Force -ErrorAction Stop).Attributes -band [IO.FileAttributes]::ReparsePoint) {
        throw "Refused reparse-point directory: $Path"
    }
}

function New-SecureInstallRoot {
    $commonData = [Environment]::GetFolderPath([Environment+SpecialFolder]::CommonApplicationData)
    if ([string]::IsNullOrWhiteSpace($commonData)) {
        throw "Could not resolve the Windows common application-data folder."
    }

    $commonData = [IO.Path]::GetFullPath($commonData)
    if ($commonData.TrimEnd("\") -eq [IO.Path]::GetPathRoot($commonData).TrimEnd("\")) {
        throw "Refused to use a drive root for W11 Optimiser staging."
    }

    $productRoot = Join-Path $commonData "W11 Optimiser"
    $stagingRoot = Join-Path $productRoot "Staging"
    foreach ($path in @($productRoot, $stagingRoot)) {
        if (Test-Path -LiteralPath $path) {
            Assert-DirectoryIsNotReparsePoint -Path $path
            Assert-AdminOnlyPathAcl -Path $path
        }
        else {
            [void] (New-Item -ItemType Directory -Path $path -ErrorAction Stop)
            Set-AdminOnlyDirectoryAcl -Path $path
            Assert-AdminOnlyPathAcl -Path $path
        }
    }

    $installRoot = Join-Path $stagingRoot ([Guid]::NewGuid().ToString("N"))
    [void] (New-Item -ItemType Directory -Path $installRoot -ErrorAction Stop)
    Set-AdminOnlyDirectoryAcl -Path $installRoot
    Assert-AdminOnlyPathAcl -Path $installRoot
    return $installRoot
}

function Invoke-ElevatedBootstrap {
    Assert-HttpsReleaseUri -Uri $bootstrapUrl
    $initiatingUserSid = Get-CurrentUserSid
    $command = @"
`$ErrorActionPreference = 'Stop'
[Net.ServicePointManager]::SecurityProtocol = [Net.SecurityProtocolType]::Tls12
`$response = Invoke-WebRequest -UseBasicParsing -MaximumRedirection 0 -Uri '$bootstrapUrl'
`$bootstrapScript = [ScriptBlock]::Create(`$response.Content)
& `$bootstrapScript -ExpectedUserSid '$initiatingUserSid'
"@
    $encoded = [Convert]::ToBase64String([Text.Encoding]::Unicode.GetBytes($command))
    $process = Start-Process -FilePath "powershell.exe" -ArgumentList @(
        "-NoProfile", "-ExecutionPolicy", "Bypass", "-EncodedCommand", $encoded
    ) -Verb RunAs -Wait -PassThru -ErrorAction Stop
    if ($null -ne $process.ExitCode -and $process.ExitCode -ne 0) {
        throw "Elevated W11 Optimiser bootstrap exited with code $($process.ExitCode)."
    }
}

function Assert-StagedPayload {
    if ([string]::IsNullOrWhiteSpace($script:ScriptPath) -or
        [string]::IsNullOrWhiteSpace($script:ExpectedScriptHash)) {
        throw "The verified W11 Optimiser payload is not available."
    }

    $root = [IO.Path]::GetFullPath($script:InstallRoot).TrimEnd("\") + "\"
    $payload = [IO.Path]::GetFullPath($script:ScriptPath)
    if (-not $payload.StartsWith($root, [StringComparison]::OrdinalIgnoreCase)) {
        throw "Refused payload outside the secure staging directory."
    }
    Assert-DirectoryIsNotReparsePoint -Path $script:InstallRoot
    Assert-AdminOnlyPathAcl -Path $script:InstallRoot
    Assert-AdminOnlyPathAcl -Path $payload

    $actualHash = (Get-FileHash -LiteralPath $payload -Algorithm SHA256 -ErrorAction Stop).Hash
    if ($actualHash -ne $script:ExpectedScriptHash) {
        throw "The staged W11 Optimiser payload changed after verification. Nothing was run."
    }
}

function Write-Title {
    Clear-Host
    Write-Host ""
    Write-Host "==========================================================================" -ForegroundColor Cyan
    Write-Host ""
    Write-Host "█   █   █     █       ███  ████  █████ ███ █   █ ███  ████ █████ ████  " -ForegroundColor White
    Write-Host "█   █  ██    ██      █   █ █   █   █    █  ██ ██  █  █     █     █   █ " -ForegroundColor White
    Write-Host "█ █ █   █     █      █   █ ████    █    █  █ █ █  █   ███  ████  ████  " -ForegroundColor White
    Write-Host "██ ██   █     █      █   █ █       █    █  █   █  █      █ █     █  █  " -ForegroundColor White
    Write-Host "█   █  ███   ███      ███  █       █   ███ █   █ ███ ████  █████ █   █ " -ForegroundColor White
    Write-Host ""
    Write-Host "                              by Julian Baumgardt" -ForegroundColor DarkGray
    Write-Host "                              version $script:W11Version" -ForegroundColor DarkGray
    Write-Host ""
    Write-Host "==========================================================================" -ForegroundColor Cyan
    Write-Host ""
}

function Download-W11Optimiser {
    Assert-HttpsReleaseUri -Uri $manifestUrl
    Assert-HttpsReleaseUri -Uri $scriptUrl
    Write-Host "CHECKING LATEST W11 OPTIMISER RELEASE..." -ForegroundColor Cyan
    Invoke-WebRequest -UseBasicParsing -MaximumRedirection 0 -Uri $manifestUrl -OutFile $script:ManifestPath
    Set-AdminOnlyFileAcl -Path $script:ManifestPath
    $manifest = Get-Content -LiteralPath $script:ManifestPath -Raw | ConvertFrom-Json

    if ($manifest.script -ne "w11-optimiser.ps1" -or
        [string]::IsNullOrWhiteSpace($manifest.version) -or
        $manifest.sha256 -notmatch "^[a-fA-F0-9]{64}$") {
        throw "The W11 Optimiser release manifest is invalid. No script was run."
    }

    $script:W11Version = $manifest.version
    Write-Host "DOWNLOADING W11 OPTIMISER VERSION $script:W11Version..." -ForegroundColor Cyan
    Invoke-WebRequest -UseBasicParsing -MaximumRedirection 0 -Uri $scriptUrl -OutFile $script:ScriptPath
    Set-AdminOnlyFileAcl -Path $script:ScriptPath
    $actualHash = (Get-FileHash -LiteralPath $script:ScriptPath -Algorithm SHA256).Hash
    $script:ExpectedScriptHash = $manifest.sha256.ToUpperInvariant()
    if ($actualHash -ne $script:ExpectedScriptHash) {
        Remove-Item -LiteralPath $script:ScriptPath -Force -ErrorAction SilentlyContinue
        throw "Downloaded script hash did not match the release manifest. No script was run."
    }
    Write-Host "SHA-256 MATCHED RELEASE MANIFEST." -ForegroundColor Green
    try {
        Unblock-File -LiteralPath $script:ScriptPath -ErrorAction Stop
    }
    catch {
    }
}

function Invoke-W11Optimiser {
    param(
        [Parameter(Mandatory = $true)][string] $Mode,
        [string[]] $ExtraArgs = @()
    )

    Assert-StagedPayload
    $identityArgs = @("-ExpectedUserSid", (Get-CurrentUserSid))
    & powershell.exe -NoProfile -ExecutionPolicy Bypass -File $script:ScriptPath -Mode $Mode @identityArgs @ExtraArgs
    if ($null -eq $LASTEXITCODE) {
        $script:W11LastExitCode = 0
    }
    else {
        $script:W11LastExitCode = [int] $LASTEXITCODE
    }
}

function Invoke-W11MenuAction {
    param(
        [Parameter(Mandatory = $true)][string] $Mode,
        [string[]] $ExtraArgs = @()
    )

    Invoke-W11Optimiser -Mode $Mode -ExtraArgs $ExtraArgs
    $exitCode = $script:W11LastExitCode
    Write-Host ""
    if ($exitCode -eq 0) {
        Write-Host "Done. If a report was generated, it should open in your browser." -ForegroundColor Green
        Write-Host "Files are saved in Desktop\W11 Optimiser." -ForegroundColor DarkGray
    }
    else {
        Write-Host "Stopped or failed. No further menu action was taken." -ForegroundColor Yellow
        Write-Host "Exit code: $exitCode" -ForegroundColor DarkGray
    }
    Write-Host ""
    [void] (Read-Host "Press Enter To Return To The Menu")
}

if ($LoadFunctionsOnly) {
    return
}

if (-not (Test-IsAdministrator)) {
    try {
        Invoke-ElevatedBootstrap
    }
    catch {
        Write-Host "W11 Optimiser elevation failed. Nothing was run." -ForegroundColor Red
        Write-Host $_.Exception.Message -ForegroundColor DarkGray
    }
    return
}

try {
    Assert-ExpectedUserIdentity -ExpectedSid $ExpectedUserSid
}
catch {
    Write-Host "W11 Optimiser refused cross-account elevation. Nothing was run." -ForegroundColor Red
    Write-Host $_.Exception.Message -ForegroundColor DarkGray
    exit 1
}

[Net.ServicePointManager]::SecurityProtocol = [Net.SecurityProtocolType]::Tls12
$script:InstallRoot = New-SecureInstallRoot
$script:ScriptPath = Join-Path $script:InstallRoot "w11-optimiser.ps1"
$script:ManifestPath = Join-Path $script:InstallRoot "w11-optimiser.manifest.json"

try {
    Write-Title
    Write-Host "This command downloads the latest W11 Optimiser script from:" -ForegroundColor DarkGray
    Write-Host "  $scriptUrl" -ForegroundColor DarkGray
    Write-Host ""
    Write-Host "No optimisation starts until you choose an option." -ForegroundColor Yellow
    Write-Host ""
    try {
        Download-W11Optimiser
    }
    catch {
        Write-Host ""
        Write-Host "DOWNLOAD FAILED. NO CHANGES WERE MADE." -ForegroundColor Red
        Write-Host $_.Exception.Message -ForegroundColor DarkGray
        return
    }

    while ($true) {
        Write-Title
        Write-Host "Recommended" -ForegroundColor White
        Write-Host ""
        Write-Host "  X  Safe Optimise (No Temp/Cache Cleanup)"
        Write-Host ""
        Write-Host "Reports" -ForegroundColor White
        Write-Host ""
        Write-Host "  1  Preview"
        Write-Host "  2  Audit"
        Write-Host "  3  Post-Check"
        Write-Host "  4  Open Last Report"
        Write-Host ""
        Write-Host "Maintenance" -ForegroundColor White
        Write-Host ""
        Write-Host "  5  Safe Optimise + Temp/Cache Cleanup"
        Write-Host "  6  Undo Latest Run"
        Write-Host ""
        Write-Host "Other" -ForegroundColor White
        Write-Host ""
        Write-Host "  0  Exit"
        Write-Host ""

        $choice = Read-Host "Choose An Option"

        switch -Regex ($choice) {
            "^x$" { Invoke-W11MenuAction -Mode "SafeOptimize" -ExtraArgs @("-SkipTempCleanup") }
            "^1$" { Invoke-W11MenuAction -Mode "Preview" }
            "^2$" { Invoke-W11MenuAction -Mode "Audit" }
            "^3$" { Invoke-W11MenuAction -Mode "PostCheck" }
            "^4$" { Invoke-W11MenuAction -Mode "OpenLastReport" }
            "^5$" { Invoke-W11MenuAction -Mode "SafeOptimize" }
            "^6$" { Invoke-W11MenuAction -Mode "UndoLatest" }
            "^0$" {
                Write-Host "Cancelled. No changes were made." -ForegroundColor DarkGray
                return
            }
            default {
                Write-Host "Please choose X or 0-6." -ForegroundColor Yellow
                Start-Sleep -Seconds 1
            }
        }
    }
}
finally {
    if (-not [string]::IsNullOrWhiteSpace($script:InstallRoot) -and
        (Test-Path -LiteralPath $script:InstallRoot)) {
        Remove-Item -LiteralPath $script:InstallRoot -Recurse -Force -ErrorAction SilentlyContinue
    }
}
