# Requires running as administrator
# Installs Firefox ESR, configures it as the default browser and force-installs uBlock Origin.

[CmdletBinding(SupportsShouldProcess = $true)]
param(
    [string]$FirefoxPackageId = "Mozilla.Firefox.ESR.fr",
    [string]$AssociationsPath = "$env:ProgramData\WindowsLibreSoftwareToolkit\FirefoxDefaultAssociations.xml"
)

$ErrorActionPreference = "Stop"

. (Join-Path $PSScriptRoot "Initialize-Winget.ps1")

$IsAdministrator = ([Security.Principal.WindowsPrincipal] [Security.Principal.WindowsIdentity]::GetCurrent()).IsInRole(
    [Security.Principal.WindowsBuiltInRole]::Administrator
)

if (-not $IsAdministrator) {
    Write-Error "This script must be run from an administrator PowerShell session."
    exit 1
}

function Install-FirefoxEsr {
    param(
        [string]$PackageId
    )

    Initialize-Winget

    Write-Host "Installing Firefox ESR with winget: $PackageId" -ForegroundColor Cyan

    if ($PSCmdlet.ShouldProcess($PackageId, "Install Firefox ESR")) {
        $Arguments = @(
            "install",
            "--id", $PackageId,
            "--silent",
            "--accept-source-agreements",
            "--accept-package-agreements"
        )

        $Process = Start-Process -FilePath "winget" -ArgumentList $Arguments -NoNewWindow -Wait -PassThru

        if ($Process.ExitCode -ne 0) {
            Write-Error "Firefox ESR installation failed. winget exit code: $($Process.ExitCode)"
            exit $Process.ExitCode
        }

        Write-Host "Firefox ESR installed or already present." -ForegroundColor Green
    }
}

function Get-FirefoxExecutablePath {
    $FirefoxPaths = Get-FirefoxExecutablePaths

    if ($FirefoxPaths.Count -gt 0) {
        return $FirefoxPaths[0]
    }

    return $null
}

function Get-FirefoxExecutablePaths {
    $CandidatePaths = @()

    $AppPath = Get-ItemProperty -Path "HKLM:\SOFTWARE\Microsoft\Windows\CurrentVersion\App Paths\firefox.exe" -ErrorAction SilentlyContinue
    if ($AppPath -and $AppPath.'(default)') {
        $CandidatePaths += $AppPath.'(default)'
    }

    $UninstallRoots = @(
        "HKLM:\SOFTWARE\Microsoft\Windows\CurrentVersion\Uninstall\*",
        "HKLM:\SOFTWARE\WOW6432Node\Microsoft\Windows\CurrentVersion\Uninstall\*"
    )

    foreach ($Root in $UninstallRoots) {
        $FirefoxInstall = Get-ItemProperty -Path $Root -ErrorAction SilentlyContinue |
            Where-Object { $_.DisplayName -like "Mozilla Firefox*" -and $_.InstallLocation } |
            Select-Object -First 1

        if ($FirefoxInstall) {
            $CandidatePaths += (Join-Path $FirefoxInstall.InstallLocation "firefox.exe")
        }
    }

    $CandidatePaths += @(
        "$env:ProgramFiles\Mozilla Firefox\firefox.exe",
        "$env:ProgramFiles\Mozilla Firefox ESR\firefox.exe",
        "${env:ProgramFiles(x86)}\Mozilla Firefox\firefox.exe"
        "${env:ProgramFiles(x86)}\Mozilla Firefox ESR\firefox.exe"
    )

    $FirefoxPaths = @()

    foreach ($Path in ($CandidatePaths | Where-Object { -not [string]::IsNullOrWhiteSpace($_) } | Select-Object -Unique)) {
        if (Test-Path $Path) {
            $FirefoxPaths += (Resolve-Path $Path).Path
        }
    }

    return @($FirefoxPaths | Select-Object -Unique)
}

function Get-FirefoxAssociationProgIds {
    $CapabilitiesPath = "HKLM:\SOFTWARE\Clients\StartMenuInternet\FIREFOX.EXE\Capabilities"
    $UrlAssociationsPath = Join-Path $CapabilitiesPath "URLAssociations"
    $FileAssociationsPath = Join-Path $CapabilitiesPath "FileAssociations"

    $UrlAssociations = Get-ItemProperty -Path $UrlAssociationsPath -ErrorAction SilentlyContinue
    $FileAssociations = Get-ItemProperty -Path $FileAssociationsPath -ErrorAction SilentlyContinue

    $HttpProgId = $UrlAssociations.http
    $HttpsProgId = $UrlAssociations.https
    $HtmlProgId = $FileAssociations.".html"
    $HtmProgId = $FileAssociations.".htm"

    if (-not $HttpProgId) {
        $HttpProgId = "FirefoxURL-308046B0AF4A39CB"
    }

    if (-not $HttpsProgId) {
        $HttpsProgId = $HttpProgId
    }

    if (-not $HtmlProgId) {
        $HtmlProgId = "FirefoxHTML-308046B0AF4A39CB"
    }

    if (-not $HtmProgId) {
        $HtmProgId = $HtmlProgId
    }

    return [PSCustomObject]@{
        Http  = $HttpProgId
        Https = $HttpsProgId
        Html  = $HtmlProgId
        Htm   = $HtmProgId
    }
}

function Get-FirefoxEnterprisePolicies {
    $Policies = [ordered]@{
        policies = [ordered]@{
            DontCheckDefaultBrowser = $true
            DisableDefaultBrowserAgent = $true
            FirefoxHome = [ordered]@{
                Search             = $true
                TopSites           = $false
                SponsoredTopSites  = $false
                Pocket             = $false
                SponsoredPocket    = $false
            }
            ExtensionSettings = [ordered]@{
                "uBlock0@raymondhill.net" = [ordered]@{
                    installation_mode = "force_installed"
                    install_url = "https://addons.mozilla.org/firefox/downloads/latest/ublock-origin/latest.xpi"
                }
            }
            "3rdparty" = [ordered]@{
                Extensions = [ordered]@{
                    "uBlock0@raymondhill.net" = [ordered]@{
                        adminSettings = [ordered]@{
                            selectedFilterLists = @(
                                "user-filters",
                                "ublock-filters",
                                "ublock-badware",
                                "ublock-privacy",
                                "ublock-abuse",
                                "ublock-unbreak",
                                "ublock-quick-fixes",
                                "easylist",
                                "easyprivacy",
                                "urlhaus-1",
                                "plowe-0"
                            )
                            userSettings = [ordered]@{
                                advancedUserEnabled = $false
                                cloudStorageEnabled = $false
                                contextMenuEnabled = $true
                                showIconBadge = $true
                            }
                        }
                    }
                }
            }
        }
    }

    return $Policies
}

function Set-FirefoxEnterprisePolicies {
    param(
        [string[]]$FirefoxExecutables
    )

    $Policies = Get-FirefoxEnterprisePolicies
    $PoliciesJson = $Policies | ConvertTo-Json -Depth 12
    $Utf8NoBom = [System.Text.UTF8Encoding]::new($false)

    foreach ($FirefoxExe in ($FirefoxExecutables | Select-Object -Unique)) {
        $FirefoxDirectory = Split-Path -Parent $FirefoxExe
        $DistributionDirectory = Join-Path $FirefoxDirectory "distribution"
        $PoliciesPath = Join-Path $DistributionDirectory "policies.json"

        if ($PSCmdlet.ShouldProcess($PoliciesPath, "Write Firefox enterprise policies")) {
            New-Item -Path $DistributionDirectory -ItemType Directory -Force | Out-Null
            [System.IO.File]::WriteAllText($PoliciesPath, $PoliciesJson, $Utf8NoBom)
            Get-Content -Path $PoliciesPath -Raw | ConvertFrom-Json | Out-Null
            Write-Host "Firefox enterprise policies written: $PoliciesPath" -ForegroundColor Green
        }
    }
}

function Set-FirefoxDefaultAssociations {
    param(
        [string]$OutputPath
    )

    $ProgIds = Get-FirefoxAssociationProgIds
    $OutputDirectory = Split-Path -Parent $OutputPath

    $Xml = @"
<?xml version="1.0" encoding="UTF-8"?>
<DefaultAssociations>
  <Association Identifier=".htm" ProgId="$($ProgIds.Htm)" ApplicationName="Firefox" />
  <Association Identifier=".html" ProgId="$($ProgIds.Html)" ApplicationName="Firefox" />
  <Association Identifier="http" ProgId="$($ProgIds.Http)" ApplicationName="Firefox" />
  <Association Identifier="https" ProgId="$($ProgIds.Https)" ApplicationName="Firefox" />
</DefaultAssociations>
"@

    if ($PSCmdlet.ShouldProcess($OutputPath, "Write Firefox default app associations")) {
        New-Item -Path $OutputDirectory -ItemType Directory -Force | Out-Null
        Set-Content -Path $OutputPath -Value $Xml -Encoding UTF8
        Write-Host "Default app associations written: $OutputPath" -ForegroundColor Green
    }

    $SystemPolicyPath = "HKLM:\SOFTWARE\Policies\Microsoft\Windows\System"

    if ($PSCmdlet.ShouldProcess($SystemPolicyPath, "Set default associations policy")) {
        New-Item -Path $SystemPolicyPath -Force | Out-Null
        New-ItemProperty -Path $SystemPolicyPath -Name "DefaultAssociationsConfiguration" -Value $OutputPath -PropertyType String -Force | Out-Null
        Write-Host "Default associations policy configured." -ForegroundColor Green
    }

    if ($PSCmdlet.ShouldProcess("DISM", "Import default app associations for new users")) {
        $Arguments = "/Online", "/Import-DefaultAppAssociations:$OutputPath"
        $Process = Start-Process -FilePath "dism.exe" -ArgumentList $Arguments -NoNewWindow -Wait -PassThru

        if ($Process.ExitCode -ne 0) {
            Write-Warning "DISM failed to import default app associations. Exit code: $($Process.ExitCode)"
        }
        else {
            Write-Host "Default app associations imported for future users." -ForegroundColor Green
        }
    }
}

Install-FirefoxEsr -PackageId $FirefoxPackageId

$FirefoxExe = Get-FirefoxExecutablePath
$FirefoxExecutables = Get-FirefoxExecutablePaths

if (-not $FirefoxExe -or $FirefoxExecutables.Count -eq 0) {
    Write-Error "Firefox was not found after installation. Check the winget package result and retry."
    exit 1
}

Write-Host "Firefox found: $FirefoxExe" -ForegroundColor Cyan
Set-FirefoxEnterprisePolicies -FirefoxExecutables $FirefoxExecutables
Set-FirefoxDefaultAssociations -OutputPath $AssociationsPath

Write-Host ""
Write-Host "Firefox ESR installed and configured. Restart Windows or sign out before creating/testing new users." -ForegroundColor Green