using namespace System.Xml;

$null = New-Module {
    $associationElementName = "Association";
    $rootSelector = "/DefaultAssociations";
    $associationSelector = "./$associationElementName";

    <#
        .SYNOPSIS
        Generates a selector for getting the association of the specified identifier.

        .PARAMETER Identifier
        The identifier to get the association of.
    #>
    $getSelector = {
        param(
            [string] $Identifier
        )

        [string] $filter = "";

        if ($Identifier) {
            $filter = "[@Identifier='$Identifier']";
        }

        "$rootSelector/$associationSelector$filter";
    };

    <#
        .SYNOPSIS
        Gets the default app associations of the running Windows system.
    #>
    function Get-DefaultAppAssociations {
        [OutputType([xml])]
        param()
        $result = [xml]::new();
        $result.LoadXml(((DISM /Online /Get-DefaultAppAssociations) | Select-Object -Skip 6 | Select-Object -SkipLast 2 | Out-String));
        $result;
    }

    <#
        .SYNOPSIS
        Sets a default app association.

        .PARAMETER Identifier
        The identifier of the association to set.

        .PARAMETER ProgId
        The ProgId to set the association to.

        .PARAMETER ApplicationName
        The ApplicationName to set the association to.
    #>
    function Set-DefaultAppAssociation {
        param(
            [string] $Identifier,
            [string] $ProgId,
            [string] $ApplicationName
        )

        [System.Xml.XmlNode] $association = $null;
        $document = Get-DefaultAppAssociations;
        $candidates = $document.SelectNodes((& $getSelector $Identifier));

        if ($candidates.Count -eq 1) {
            $association = $candidates[0];
        }
        else {
            $association = $document.SelectSingleNode($rootSelector).AppendChild($document.CreateElement($associationElementName));

            foreach ($attributeName in @("Identifier", "ProgId", "ApplicationName")) {
                $null = $association.Attributes.Append($document.CreateAttribute($attributeName));
            }

            $association.Identifier = $Identifier;
        }

        $association.ProgId = $ProgId;
        $association.ApplicationName = $ApplicationName;
        Save-DefaultAppAssociations $document;
    }

    <#
        .SYNOPSIS
        Saves the default app associations.

        .PARAMETER Document
        The xml document containing the declarations of the app associations.
    #>
    function Save-DefaultAppAssociations {
        param(
            [xml] $Document
        )

        $root = $document.SelectSingleNode($rootSelector);
        $associations = $root.SelectNodes((& $getSelector));

        # Reorder associations by their Identifier
        $null = $associations | ForEach-Object { $root.RemoveChild($_) } | Sort-Object -Property "Identifier" | ForEach-Object { $root.AppendChild($_) };

        $configFile = New-TemporaryFile;
        $writerSettings = [XmlWriterSettings]::new();
        $writerSettings.Indent = $true;
        $writerSettings.Encoding = [System.Text.UTF8Encoding]::new();
        $writer = [XmlWriter]::Create($configFile.FullName, $writerSettings);
        $document.Save($writer);
        $writer.Dispose();
        $null = DISM /Online "/Import-DefaultAppAssociations:$($configFile.FullName)";
        Remove-Item $configFile;
    }
}