[CmdletBinding()]
param(
    [string]$ExePath = (Join-Path $PSScriptRoot 'dist\PortableCaffeine.exe'),
    [string]$Subject = 'CN=Portable Caffeine Dev Code Signing',
    [string]$FriendlyName = 'Portable Caffeine Dev Code Signing (Self-Signed)',
    [ValidateSet('None', 'CurrentUser', 'LocalMachine')]
    [string]$TrustScope = 'CurrentUser',
    [ValidateSet('CurrentUser', 'LocalMachine')]
    [string]$PrivateKeyStoreScope = 'CurrentUser',
    [int]$YearsValid = 5,
    [switch]$ForceNewCertificate,
    [string]$TimestampServer,
    [string]$CerExportPath = (Join-Path $PSScriptRoot 'dist\PortableCaffeine-dev-signing.cer'),
    [string]$PfxExportPath,
    [securestring]$PfxPassword
)

Set-StrictMode -Version Latest
$ErrorActionPreference = 'Stop'

if ([Environment]::OSVersion.Platform -ne [PlatformID]::Win32NT) {
    throw "This script is intended for Windows only."
}

if (-not (Get-Command New-SelfSignedCertificate -ErrorAction SilentlyContinue)) {
    throw "New-SelfSignedCertificate is not available in this PowerShell session."
}

if (-not (Get-Command Set-AuthenticodeSignature -ErrorAction SilentlyContinue)) {
    throw "Set-AuthenticodeSignature is not available in this PowerShell session."
}

if (-not ('PortableCaffeineBuild.SigningHelpers' -as [type])) {
    $helperCode = @"
using System;
using System.Security.Principal;
using System.Security.Cryptography.X509Certificates;

namespace PortableCaffeineBuild
{
    public static class SigningHelpers
    {
        public static bool IsAdministrator()
        {
            WindowsIdentity identity = WindowsIdentity.GetCurrent();
            if (identity == null)
            {
                return false;
            }

            WindowsPrincipal principal = new WindowsPrincipal(identity);
            return principal.IsInRole(WindowsBuiltInRole.Administrator);
        }

        public static bool StoreContainsThumbprint(string storeName, string storeLocation, string thumbprint)
        {
            if (string.IsNullOrEmpty(thumbprint))
            {
                return false;
            }

            StoreName sn = (StoreName)Enum.Parse(typeof(StoreName), storeName, true);
            StoreLocation sl = (StoreLocation)Enum.Parse(typeof(StoreLocation), storeLocation, true);

            X509Store store = new X509Store(sn, sl);
            try
            {
                store.Open(OpenFlags.ReadOnly);
                string normalized = thumbprint.Replace(" ", string.Empty).ToUpperInvariant();
                foreach (X509Certificate2 cert in store.Certificates)
                {
                    if (cert == null || string.IsNullOrEmpty(cert.Thumbprint))
                    {
                        continue;
                    }

                    if (cert.Thumbprint.Replace(" ", string.Empty).ToUpperInvariant() == normalized)
                    {
                        return true;
                    }
                }

                return false;
            }
            finally
            {
                store.Close();
            }
        }

        public static void AddPublicCertificateToStore(byte[] rawData, string storeName, string storeLocation)
        {
            if (rawData == null || rawData.Length == 0)
            {
                throw new ArgumentException("rawData");
            }

            X509Certificate2 publicCert = new X509Certificate2(rawData);
            try
            {
                if (StoreContainsThumbprint(storeName, storeLocation, publicCert.Thumbprint))
                {
                    return;
                }

                StoreName sn = (StoreName)Enum.Parse(typeof(StoreName), storeName, true);
                StoreLocation sl = (StoreLocation)Enum.Parse(typeof(StoreLocation), storeLocation, true);
                X509Store store = new X509Store(sn, sl);
                try
                {
                    store.Open(OpenFlags.ReadWrite);
                    store.Add(publicCert);
                }
                finally
                {
                    store.Close();
                }
            }
            finally
            {
                publicCert.Reset();
            }
        }
    }
}
"@

    Add-Type -TypeDefinition $helperCode -ReferencedAssemblies @(
        'System.dll'
    )
}

function Test-IsAdministrator {
    return [PortableCaffeineBuild.SigningHelpers]::IsAdministrator()
}

function Get-CodeSigningCertificate {
    param(
        [Parameter(Mandatory = $true)]
        [string]$SubjectName,
        [Parameter(Mandatory = $true)]
        [string]$StoreScope
    )

    $storePath = "Cert:\$StoreScope\My"
    if (-not (Test-Path -LiteralPath $storePath)) {
        return $null
    }

    $now = Get-Date
    $certs = Get-ChildItem -Path $storePath |
        Where-Object {
            $_.Subject -eq $SubjectName -and
            $_.HasPrivateKey -and
            $_.NotAfter -gt $now
        } |
        Sort-Object -Property NotAfter -Descending

    if ($certs) {
        return @($certs)[0]
    }

    return $null
}

function New-CodeSigningCertificate {
    param(
        [Parameter(Mandatory = $true)]
        [string]$SubjectName,
        [Parameter(Mandatory = $true)]
        [string]$Friendly,
        [Parameter(Mandatory = $true)]
        [string]$StoreScope,
        [Parameter(Mandatory = $true)]
        [int]$ValidYears
    )

    $certStoreLocation = "Cert:\$StoreScope\My"
    $notAfter = (Get-Date).AddYears($ValidYears)

    Write-Host "Creating new self-signed code signing certificate..."
    Write-Host "Subject: $SubjectName"
    Write-Host "Store:   $certStoreLocation"

    $newCert = New-SelfSignedCertificate `
        -Type CodeSigningCert `
        -Subject $SubjectName `
        -FriendlyName $Friendly `
        -CertStoreLocation $certStoreLocation `
        -KeyAlgorithm RSA `
        -KeyLength 2048 `
        -KeySpec Signature `
        -HashAlgorithm SHA256 `
        -KeyExportPolicy Exportable `
        -NotAfter $notAfter

    if (-not $newCert) {
        throw "New-SelfSignedCertificate did not return a certificate."
    }

    return $newCert
}

function Install-CertificateTrust {
    param(
        [Parameter(Mandatory = $true)]
        [System.Security.Cryptography.X509Certificates.X509Certificate2]$Certificate,
        [Parameter(Mandatory = $true)]
        [ValidateSet('CurrentUser', 'LocalMachine')]
        [string]$Scope
    )

    foreach ($storeName in @('Root', 'TrustedPublisher')) {
        $alreadyTrusted = [PortableCaffeineBuild.SigningHelpers]::StoreContainsThumbprint($storeName, $Scope, $Certificate.Thumbprint)
        if ($alreadyTrusted) {
            Write-Host "Trusted store already contains cert: $Scope\$storeName ($($Certificate.Thumbprint))"
            continue
        }

        Write-Host "Adding cert to trusted store: $Scope\$storeName"
        [PortableCaffeineBuild.SigningHelpers]::AddPublicCertificateToStore($Certificate.RawData, $storeName, $Scope)
    }
}

function Export-PublicCertificateIfRequested {
    param(
        [Parameter(Mandatory = $true)]
        [System.Security.Cryptography.X509Certificates.X509Certificate2]$Certificate,
        [string]$Path
    )

    if ([string]::IsNullOrWhiteSpace($Path)) {
        return
    }

    $dir = Split-Path -Parent $Path
    if ($dir) {
        New-Item -ItemType Directory -Force -Path $dir | Out-Null
    }

    $raw = $Certificate.Export([System.Security.Cryptography.X509Certificates.X509ContentType]::Cert)
    [System.IO.File]::WriteAllBytes($Path, $raw)
    Write-Host "Exported public certificate: $Path"
}

function Export-PfxIfRequested {
    param(
        [Parameter(Mandatory = $true)]
        [System.Security.Cryptography.X509Certificates.X509Certificate2]$Certificate,
        [string]$Path,
        [securestring]$Password
    )

    if ([string]::IsNullOrWhiteSpace($Path)) {
        return
    }

    if (-not $Password) {
        throw "PfxExportPath was provided but no -PfxPassword was supplied."
    }

    if (-not (Get-Command Export-PfxCertificate -ErrorAction SilentlyContinue)) {
        throw "Export-PfxCertificate is not available in this PowerShell session."
    }

    $dir = Split-Path -Parent $Path
    if ($dir) {
        New-Item -ItemType Directory -Force -Path $dir | Out-Null
    }

    Export-PfxCertificate -Cert $Certificate -FilePath $Path -Password $Password -Force | Out-Null
    Write-Host "Exported PFX: $Path"
}

function Sign-Executable {
    param(
        [Parameter(Mandatory = $true)]
        [string]$Path,
        [Parameter(Mandatory = $true)]
        [System.Security.Cryptography.X509Certificates.X509Certificate2]$Certificate,
        [string]$TimestampUrl
    )

    $params = @{
        FilePath      = $Path
        Certificate   = $Certificate
        HashAlgorithm = 'SHA256'
    }

    if (-not [string]::IsNullOrWhiteSpace($TimestampUrl)) {
        $params['TimestampServer'] = $TimestampUrl
    }

    Write-Host "Signing file: $Path"
    $result = Set-AuthenticodeSignature @params

    if (-not $result) {
        throw "Set-AuthenticodeSignature did not return a result."
    }

    Write-Host "Sign status: $($result.Status)"
    if ($result.StatusMessage) {
        Write-Host "Status msg: $($result.StatusMessage)"
    }

    return $result
}

if (-not (Test-Path -LiteralPath $ExePath)) {
    throw "EXE not found: $ExePath`nBuild it first with .\build-caffeine-exe.ps1"
}

if ($YearsValid -lt 1 -or $YearsValid -gt 50) {
    throw "YearsValid must be between 1 and 50."
}

$needsAdmin = $false
if ($TrustScope -eq 'LocalMachine' -or $PrivateKeyStoreScope -eq 'LocalMachine') {
    $needsAdmin = $true
}

if ($needsAdmin -and -not (Test-IsAdministrator)) {
    throw "This operation requires an elevated PowerShell session (Run as Administrator) because LocalMachine certificate stores are being modified/used."
}

$cert = $null
if (-not $ForceNewCertificate) {
    $cert = Get-CodeSigningCertificate -SubjectName $Subject -StoreScope $PrivateKeyStoreScope
}

if (-not $cert) {
    $cert = New-CodeSigningCertificate -SubjectName $Subject -Friendly $FriendlyName -StoreScope $PrivateKeyStoreScope -ValidYears $YearsValid
} else {
    Write-Host "Reusing existing certificate:"
    Write-Host "Subject:     $($cert.Subject)"
    Write-Host "Thumbprint:  $($cert.Thumbprint)"
    Write-Host "Expires:     $($cert.NotAfter)"
    Write-Host "Store:       Cert:\$PrivateKeyStoreScope\My"
}

Export-PublicCertificateIfRequested -Certificate $cert -Path $CerExportPath
Export-PfxIfRequested -Certificate $cert -Path $PfxExportPath -Password $PfxPassword

if ($TrustScope -ne 'None') {
    Write-Host ""
    Write-Host "Installing trust for self-signed certificate..."
    Install-CertificateTrust -Certificate $cert -Scope $TrustScope
} else {
    Write-Host "Skipping trust-store installation (TrustScope=None)."
}

Write-Host ""
$signResult = Sign-Executable -Path $ExePath -Certificate $cert -TimestampUrl $TimestampServer
$verify = Get-AuthenticodeSignature -FilePath $ExePath

Write-Host ""
Write-Host "Verification:"
Write-Host "Status:      $($verify.Status)"
if ($verify.SignerCertificate) {
    Write-Host "Thumbprint:  $($verify.SignerCertificate.Thumbprint)"
    Write-Host "Subject:     $($verify.SignerCertificate.Subject)"
}
if ($verify.TimeStamperCertificate) {
    Write-Host "Timestamped: Yes"
} else {
    Write-Host "Timestamped: No"
}

if ($verify.Status -ne 'Valid') {
    Write-Warning "Signature was applied but verification is not 'Valid'. For self-signed certs, ensure trust was installed in Root + TrustedPublisher for the relevant scope."
}

Write-Host ""
Write-Host "Done."
Write-Host "EXE:  $ExePath"
if ($CerExportPath) {
    Write-Host "CER:  $CerExportPath"
}