#!/bin/pwsh
$env:WIN_COMPUTER_NAME = "DerGeret";
$env:SETUP_SCRIPT_NAME = "$PSScriptRoot/Restore.ps1";

function Initialize-SetupConfig() {
    param(
        [xml] $config,
        [System.Xml.XmlNamespaceManager] $namespace
    );

    $setupComponent = $config.SelectSingleNode(
        "/ua:unattend/ua:settings[@pass='windowsPE']/ua:component[@name='Microsoft-Windows-Setup']",
        $namespace);

    $diskConfig = $setupComponent.SelectSingleNode("./ua:DiskConfiguration/ua:Disk", $namespace);

    $partitionCreationContainer = $diskConfig.SelectSingleNode("./ua:CreatePartitions", $namespace);
    $partitionCreations = $partitionCreationContainer.SelectNodes("./ua:CreatePartition", $namespace);

    $partitionModificationContainer = $diskConfig.SelectSingleNode("./ua:ModifyPartitions", $namespace);
    $partitionModifications = $partitionModificationContainer.SelectNodes("./ua:ModifyPartition", $namespace);

    <#
    .SYNOPSIS
        Gets the XML element describing the installation partition ID.
    #>
    function Get-InstallationPartition {
        $setupComponent.SelectSingleNode("./ua:ImageInstall/ua:OSImage/ua:InstallTo/ua:PartitionID", $namespace)
    }

    <#
    .SYNOPSIS
        Increases the ID of all partitions in the specified range by 1.
    #>
    function Move-PartitionRange {
        param (
            [int]$From = 0,
            [System.Nullable[int]]$To = $null,
            [int]$By = 1
        )

        # Update installation partition ID if necessary
        $installationPartition = Get-InstallationPartition;
        $installPartitionID = [int]$installationPartition.InnerText;

        if (($installPartitionID -ge $From) -and (($null -eq $To) -or ($installPartitionID -le $To))) {
            $installationPartition.InnerText = "$($installPartitionID + $By)";
        }

        # Update IDs of all partition creations
        foreach ($partition in $partitionCreations) {
            $orderNode = $partition.SelectSingleNode("./ua:Order", $namespace);
            $order = [int]$orderNode.InnerText;
            $newOrder = $order;

            if (($newOrder -ge $From) -and (($null -eq $To) -or ($newOrder -le $To))) {
                $newOrder += $By;
            }

            if ($order -ne $newOrder) {
                $orderNode.InnerText = "$newOrder";
            }
        }

        # Update IDs of all partition modifications
        foreach ($partition in $partitionModifications) {
            $partitionNode = $partition.SelectSingleNode("./ua:PartitionID", $namespace);
            $partitionID = [int]$partitionNode.InnerText;
            $newID = $partitionID;

            if (($newID -ge $From) -and (($null -eq $To) -or ($newID -le $To))) {
                $newID += $By;
            }

            if ($partitionID -ne $newID) {
                $partitionNode.InnerText = "$newID";
                $partition.SelectSingleNode("./ua:Order", $namespace).InnerText = "$newID";
            }
        }
    }

    function Add-Partition {
        param (
            [int]$Index,
            [int]$Size,
            [string]$Type = "Primary"
        )

        Move-PartitionRange -From $Index -By 1;

        $newPartition = $partitionCreations[0].CloneNode($true);
        $newPartition.SelectSingleNode("./ua:Order", $namespace).InnerText = "$Index";
        $newPartition.SelectSingleNode("./ua:Type", $namespace).InnerText = "$Type";
        $newPartition.SelectSingleNode("./ua:Size", $namespace).InnerText = "$Size";
        $null = $partitionCreationContainer.AppendChild($newPartition);

        $newModification = $partitionModifications[2].CloneNode($true);
        $newModification.SelectSingleNode("./ua:Order", $namespace).InnerText = "$Index";
        $newModification.SelectSingleNode("./ua:PartitionID", $namespace).InnerText = "$Index";
        $null = $partitionModificationContainer.AppendChild($newModification);
    }

    <#
    .SYNOPSIS
        Relocates the partition with the specified `$From` ID to the specified `$To` ID.
    #>
    function Invoke-PartitionRelocation {
        param (
            [int]$From,
            [int]$To
        )

        Move-PartitionRange $From $From (-1 * ($From + 1))

        if ($From -gt $To) {
            Move-PartitionRange $To ($From - 1);
        }
        elseif ($From -lt $To) {
            Move-PartitionRange ($From + 1) $To -1;
        }

        Move-PartitionRange -1 -1 ($To + 1)
    }

    # Resize EFI partition to 1GB
    $partitionCreations[1].SelectSingleNode("./ua:Size", $namespace).InnerText = "$(1024)";

    # Swap Windows RE partition (partition #1)  and boot partition (partition #2)
    Invoke-PartitionRelocation 2 1;

    # Add space before Windows installation... wha-!? For Linux, ofc! I use Arch Linux, btw.
    $swapSpacer = 100;
    Add-Partition 2 $swapSpacer;

    # Add a 1.2 TB partition for Linux
    Add-Partition 3 ((1.2 * 1024 * 1024) - 1024 - $swapSpacer);
}

. "$PSScriptRoot/../../../scripts/Windows/OS/Setup.ps1";