#Requires -Version 5.1

[CmdletBinding()]
param()

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

$expectedTypeLibId = '{2B7F2D64-3A8D-401C-BE73-5C0747BA342C}'
$expectedTypeLibVersion = '1.0'
$expectedThreadingModel = 'Apartment'
$amd64Machine = 0x8664
$expectedClasses = @(
    [pscustomobject]@{
        Name = 'KAEngine'
        ClassId = '{D756CDBE-AA31-42B2-9CC7-018753CA61BF}'
        ProgId = 'K3DAsyncEngine.KAEngine.1'
    },
    [pscustomobject]@{
        Name = 'KAEventHandler'
        ClassId = '{39828C77-EFF0-4E59-979B-8673C028C718}'
        ProgId = 'K3DAsyncEngine.KAEventHandler.1'
    }
)

function Get-RegistryValue {
    param(
        [Parameter(Mandatory = $true)]
        [Microsoft.Win32.RegistryKey] $BaseKey,

        [Parameter(Mandatory = $true)]
        [string] $SubKey,

        [string] $Name = ''
    )

    $key = $BaseKey.OpenSubKey($SubKey, $false)
    if ($null -eq $key) {
        return $null
    }

    try {
        return $key.GetValue(
            $Name,
            $null,
            [Microsoft.Win32.RegistryValueOptions]::DoNotExpandEnvironmentNames)
    }
    finally {
        $key.Dispose()
    }
}

function Assert-TextEqual {
    param(
        [AllowNull()]
        [object] $Actual,

        [Parameter(Mandatory = $true)]
        [string] $Expected,

        [Parameter(Mandatory = $true)]
        [string] $Description
    )

    if ($null -eq $Actual -or
        -not [string]::Equals(
            [string] $Actual,
            $Expected,
            [StringComparison]::OrdinalIgnoreCase)) {
        throw "$Description mismatch. Expected '$Expected'; found '$Actual'."
    }
}

function Assert-GuidEqual {
    param(
        [AllowNull()]
        [object] $Actual,

        [Parameter(Mandatory = $true)]
        [string] $Expected,

        [Parameter(Mandatory = $true)]
        [string] $Description
    )

    $actualGuid = [Guid]::Empty
    $expectedGuid = [Guid]::Empty
    if ($null -eq $Actual -or
        -not [Guid]::TryParse([string] $Actual, [ref] $actualGuid) -or
        -not [Guid]::TryParse($Expected, [ref] $expectedGuid) -or
        $actualGuid -ne $expectedGuid) {
        throw "$Description mismatch. Expected '$Expected'; found '$Actual'."
    }
}

function Resolve-RegisteredFilePath {
    param(
        [AllowNull()]
        [object] $RegistryValue,

        [Parameter(Mandatory = $true)]
        [string] $Description
    )

    if ($null -eq $RegistryValue -or [string]::IsNullOrWhiteSpace([string] $RegistryValue)) {
        throw "$Description is not registered."
    }

    $expanded = [Environment]::ExpandEnvironmentVariables([string] $RegistryValue)
    $expanded = $expanded.Trim().Trim('"')
    if (-not [IO.Path]::IsPathRooted($expanded)) {
        throw "$Description is not an absolute path."
    }

    $fullPath = [IO.Path]::GetFullPath($expanded)
    if (-not (Test-Path -LiteralPath $fullPath -PathType Leaf)) {
        throw "$Description does not exist: $fullPath"
    }

    return $fullPath
}

function Get-PeMachine {
    param(
        [Parameter(Mandatory = $true)]
        [string] $Path
    )

    $stream = [IO.File]::Open(
        $Path,
        [IO.FileMode]::Open,
        [IO.FileAccess]::Read,
        ([IO.FileShare]::ReadWrite -bor [IO.FileShare]::Delete))
    $reader = [IO.BinaryReader]::new($stream)

    try {
        if ($stream.Length -lt 64 -or $reader.ReadUInt16() -ne 0x5A4D) {
            throw "Not a PE file: $Path"
        }

        $stream.Position = 0x3C
        $peOffset = $reader.ReadInt32()
        if ($peOffset -lt 0 -or ($peOffset + 6) -gt $stream.Length) {
            throw "Invalid PE header offset: $Path"
        }

        $stream.Position = $peOffset
        if ($reader.ReadUInt32() -ne 0x00004550) {
            throw "Invalid PE signature: $Path"
        }

        $machineCode = $reader.ReadUInt16()
        $machineName = switch ($machineCode) {
            0x014C { 'I386' }
            0x8664 { 'AMD64' }
            0xAA64 { 'ARM64' }
            default { 'Unknown' }
        }

        return [pscustomobject]@{
            Code = $machineCode
            Name = $machineName
        }
    }
    finally {
        $reader.Dispose()
        $stream.Dispose()
    }
}

function Assert-Amd64File {
    param(
        [Parameter(Mandatory = $true)]
        [string] $Path,

        [Parameter(Mandatory = $true)]
        [string] $Description
    )

    $machine = Get-PeMachine -Path $Path
    if ($machine.Code -ne $amd64Machine) {
        throw "$Description is not AMD64. PE machine is $($machine.Name) (0x$('{0:X4}' -f $machine.Code))."
    }
}

function Test-RegistryKeyExists {
    param(
        [Parameter(Mandatory = $true)]
        [Microsoft.Win32.RegistryKey] $BaseKey,

        [Parameter(Mandatory = $true)]
        [string] $SubKey
    )

    $key = $BaseKey.OpenSubKey($SubKey, $false)
    if ($null -eq $key) {
        return $false
    }

    $key.Dispose()
    return $true
}

function Assert-NoCurrentUserOverrides {
    param(
        [Parameter(Mandatory = $true)]
        [Microsoft.Win32.RegistryKey] $CurrentUserRoot
    )

    $paths = @("Software\Classes\TypeLib\$expectedTypeLibId")
    foreach ($registration in $expectedClasses) {
        $paths += "Software\Classes\CLSID\$($registration.ClassId)"
        $paths += "Software\Classes\$($registration.ProgId)"
    }

    foreach ($path in $paths) {
        if (Test-RegistryKeyExists -BaseKey $CurrentUserRoot -SubKey $path) {
            throw "A per-user Registry64 K3D COM override is present at HKCU:\$path. Remove the override through the approved vendor recovery procedure before use."
        }
    }
}

function Assert-ComClassRegistration {
    param(
        [Parameter(Mandatory = $true)]
        [Microsoft.Win32.RegistryKey] $MachineRoot,

        [Parameter(Mandatory = $true)]
        [object] $Registration
    )

    $classesPrefix = 'SOFTWARE\Classes'
    $progIdClassId = Get-RegistryValue `
        -BaseKey $MachineRoot `
        -SubKey "$classesPrefix\$($Registration.ProgId)\CLSID"
    Assert-GuidEqual `
        -Actual $progIdClassId `
        -Expected $Registration.ClassId `
        -Description "$($Registration.Name) ProgID to CLSID mapping"

    $classProgId = Get-RegistryValue `
        -BaseKey $MachineRoot `
        -SubKey "$classesPrefix\CLSID\$($Registration.ClassId)\ProgID"
    Assert-TextEqual `
        -Actual $classProgId `
        -Expected $Registration.ProgId `
        -Description "$($Registration.Name) CLSID to ProgID mapping"

    $classTypeLibId = Get-RegistryValue `
        -BaseKey $MachineRoot `
        -SubKey "$classesPrefix\CLSID\$($Registration.ClassId)\TypeLib"
    Assert-GuidEqual `
        -Actual $classTypeLibId `
        -Expected $expectedTypeLibId `
        -Description "$($Registration.Name) TypeLib GUID"

    $classTypeLibVersion = Get-RegistryValue `
        -BaseKey $MachineRoot `
        -SubKey "$classesPrefix\CLSID\$($Registration.ClassId)\Version"
    if ($null -ne $classTypeLibVersion -and
        -not [string]::IsNullOrWhiteSpace([string] $classTypeLibVersion)) {
        Assert-TextEqual `
            -Actual $classTypeLibVersion `
            -Expected $expectedTypeLibVersion `
            -Description "$($Registration.Name) TypeLib version"
    }

    $inprocKey = "$classesPrefix\CLSID\$($Registration.ClassId)\InprocServer32"
    $serverPath = Resolve-RegisteredFilePath `
        -RegistryValue (Get-RegistryValue -BaseKey $MachineRoot -SubKey $inprocKey) `
        -Description "$($Registration.Name) InprocServer32 path"
    $threadingModel = Get-RegistryValue `
        -BaseKey $MachineRoot `
        -SubKey $inprocKey `
        -Name 'ThreadingModel'
    Assert-TextEqual `
        -Actual $threadingModel `
        -Expected $expectedThreadingModel `
        -Description "$($Registration.Name) threading model"
    Assert-Amd64File -Path $serverPath -Description "$($Registration.Name) COM server DLL"

    return [pscustomobject][ordered]@{
        Name = $Registration.Name
        ClassId = $Registration.ClassId
        ProgId = $Registration.ProgId
        ServerPath = $serverPath
        ServerMachine = 'AMD64'
        ThreadingModel = $threadingModel
        ReciprocalMapping = $true
    }
}

if (-not [Environment]::Is64BitOperatingSystem) {
    throw 'K3D x64 registration requires a 64-bit Windows operating system.'
}

$machineRoot = [Microsoft.Win32.RegistryKey]::OpenBaseKey(
    [Microsoft.Win32.RegistryHive]::LocalMachine,
    [Microsoft.Win32.RegistryView]::Registry64)
$currentUserRoot = [Microsoft.Win32.RegistryKey]::OpenBaseKey(
    [Microsoft.Win32.RegistryHive]::CurrentUser,
    [Microsoft.Win32.RegistryView]::Registry64)

try {
    Assert-NoCurrentUserOverrides -CurrentUserRoot $currentUserRoot

    $typeLibKey = "SOFTWARE\Classes\TypeLib\$expectedTypeLibId\$expectedTypeLibVersion\0\win64"
    $typeLibPath = Resolve-RegisteredFilePath `
        -RegistryValue (Get-RegistryValue -BaseKey $machineRoot -SubKey $typeLibKey) `
        -Description 'K3D x64 type library path'
    Assert-Amd64File -Path $typeLibPath -Description 'K3D type library DLL'

    $classReports = @()
    foreach ($registration in $expectedClasses) {
        $classReports += Assert-ComClassRegistration `
            -MachineRoot $machineRoot `
            -Registration $registration
    }

    [pscustomobject][ordered]@{
        Status = 'Valid'
        RegistryHive = 'HKLM'
        RegistryView = 'Registry64'
        CurrentUserOverrides = 'None'
        TypeLibId = $expectedTypeLibId
        TypeLibVersion = $expectedTypeLibVersion
        TypeLibTarget = 'win64'
        TypeLibPath = $typeLibPath
        TypeLibMachine = 'AMD64'
        Classes = $classReports
        PowerShellProcessIs64Bit = [Environment]::Is64BitProcess
        ComActivated = $false
    }
}
finally {
    $currentUserRoot.Dispose()
    $machineRoot.Dispose()
}