if (-not ("FileAssocHelper" -as [type])) {

    Add-Type @"
using System;
using System.Text;
using System.Runtime.InteropServices;

public static class FileAssocHelper
{
    [DllImport("Shlwapi.dll", CharSet = CharSet.Unicode)]
    public static extern uint AssocQueryString(
        uint flags,
        uint str,
        string assoc,
        string extra,
        StringBuilder outString,
        ref uint outLength
    );

    public static string GetOpenCommand(string extension)
    {
        const uint ASSOCF_NONE = 0;
        const uint ASSOCSTR_COMMAND = 1;

        uint length = 4096;
        StringBuilder result = new StringBuilder((int)length);

        uint ret = AssocQueryString(
            ASSOCF_NONE,
            ASSOCSTR_COMMAND,
            extension,
            "open",
            result,
            ref length
        );

        if (ret != 0)
            return null;

        return result.ToString();
    }
}
"@

}

if (-not ("IconShellRefresh" -as [type])) {

    Add-Type @"
using System;
using System.Runtime.InteropServices;

public static class IconShellRefresh
{
    [DllImport("shell32.dll")]
    public static extern void SHChangeNotify(
        uint wEventId,
        uint uFlags,
        IntPtr dwItem1,
        IntPtr dwItem2
    );
}
"@

}

function Set-FileTypeIcon {

    param (

        [Parameter(Mandatory = $true, Position = 0)]
        [string]$Extension,

        [Parameter(Mandatory = $true, Position = 1)]
        [string]$IconPath,

        [string]$Description

    )

    if (-not $Extension.StartsWith(".")) {
        $Extension = ".$Extension"
    }

    $Extension = $Extension.ToLower()


    Write-Host ""
    Write-Host "==============================================" -ForegroundColor Cyan
    Write-Host "Configuring $Extension" -ForegroundColor Cyan
    Write-Host "=============================================="


    if (!(Test-Path $IconPath)) {

        Write-Host ""
        Write-Host "ERROR: Icon does not exist:" -ForegroundColor Red
        Write-Host $IconPath
        return
    }


    $OpenCommand =
        [FileAssocHelper]::GetOpenCommand($Extension)


    if ([string]::IsNullOrWhiteSpace($OpenCommand)) {

        Write-Host ""
        Write-Host "ERROR:" -ForegroundColor Red
        Write-Host "Windows could not determine the current application"
        Write-Host "associated with $Extension."
        Write-Host ""
        Write-Host "No changes were made."
        return
    }


    Write-Host ""
    Write-Host "Current open command:" -ForegroundColor DarkGray
    Write-Host "  $OpenCommand"


    $CleanExtension =
        $Extension.TrimStart(".")

    $ProgID =
        "Dmthie.CustomIcon.$CleanExtension"


    if ([string]::IsNullOrWhiteSpace($Description)) {

        $Description =
            "$($CleanExtension.ToUpper()) File"
    }


    Write-Host ""
    Write-Host "Custom ProgID:" -ForegroundColor DarkGray
    Write-Host "  $ProgID"


    $ProgIDPath =
        "HKCU:\Software\Classes\$ProgID"


    New-Item `
        -Path $ProgIDPath `
        -Force |
        Out-Null


    Set-Item `
        -Path $ProgIDPath `
        -Value $Description


    $IconKey =
        "$ProgIDPath\DefaultIcon"


    New-Item `
        -Path $IconKey `
        -Force |
        Out-Null


    Set-Item `
        -Path $IconKey `
        -Value "`"$IconPath`""


    $CommandKey =
        "$ProgIDPath\shell\open\command"


    New-Item `
        -Path $CommandKey `
        -Force |
        Out-Null


    Set-Item `
        -Path $CommandKey `
        -Value $OpenCommand

    $ExtensionKey =
        "HKCU:\Software\Classes\$Extension"


    New-Item `
        -Path $ExtensionKey `
        -Force |
        Out-Null


    Set-Item `
        -Path $ExtensionKey `
        -Value $ProgID

    $UserChoice =
        "HKCU:\Software\Microsoft\Windows\CurrentVersion\Explorer\FileExts\$Extension\UserChoice"


    if (Test-Path $UserChoice) {

        Write-Host ""
        Write-Host "Removing Windows UserChoice override..." `
            -ForegroundColor Yellow

        try {

            Remove-Item `
                -Path $UserChoice `
                -Recurse `
                -Force `
                -ErrorAction Stop

        }
        catch {

            Write-Host ""
            Write-Host "WARNING:" -ForegroundColor Yellow
            Write-Host "Windows prevented removal of UserChoice."
            Write-Host "The custom icon may therefore not take effect."
        }
    }


    [IconShellRefresh]::SHChangeNotify(
        0x08000000,
        0,
        [IntPtr]::Zero,
        [IntPtr]::Zero
    )


    Write-Host ""
    Write-Host "SUCCESS" -ForegroundColor Green

    Write-Host ""
    Write-Host "$Extension"
    Write-Host "Icon:"
    Write-Host "  $IconPath"

    Write-Host ""
    Write-Host "Opening behavior preserved:"
    Write-Host "  $OpenCommand"

}

function Restart-FileIconExplorer {

    Write-Host ""
    Write-Host "Restarting Windows Explorer..." `
        -ForegroundColor Cyan

    Stop-Process `
        -Name explorer `
        -Force
}