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

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

    New-Module {
        <#
        .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
            )
    
            $installPartitionID = [int]$installationPartition.InnerText;
    
            if (($installPartitionID -ge $From) -and (($null -eq $To) -or ($installPartitionID -lt $To))) {
                $installationPartition.InnerText = "$($installPartitionID + $By)";
            }
    
            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 -lt $To))) {
                    $newOrder += $By;
                }
    
                if ($order -ne $newOrder) {
                    $orderNode.InnerText = "$newOrder";
                }
            }
    
            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 -lt $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) (-1 * ($From + 1))
    
            if ($From -gt $To) {
                Move-PartitionRange $To $From;
            }
            elseif ($From -lt $To) {
                Move-PartitionRange $From ($To + 1) -1;
            }
    
            Move-PartitionRange -1 0 ($To + 1)
        }
    }

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

    $diskConfig = $setupComponent.SelectSingleNode("./ua:DiskConfiguration/ua:Disk", $namespace);
    $installationPartition = $setupComponent.SelectSingleNode("./ua:ImageInstall/ua:OSImage/ua:InstallTo/ua:PartitionID", $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);

    # 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";