#Requires -Version 5.1

[CmdletBinding(SupportsShouldProcess = $true, ConfirmImpact = 'Low')]
param(
    [string] $TlbImpPath
)

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"
}

function Resolve-X64TlbImpPath {
    param(
        [AllowEmptyString()]
        [string] $ExplicitPath
    )

    if (-not [string]::IsNullOrWhiteSpace($ExplicitPath)) {
        if (-not [IO.Path]::IsPathRooted($ExplicitPath)) {
            throw 'TlbImpPath must be an absolute path.'
        }

        $resolvedExplicitPath = [IO.Path]::GetFullPath($ExplicitPath)
        if (-not (Test-Path -LiteralPath $resolvedExplicitPath -PathType Leaf)) {
            throw "TlbImp.exe was not found: $resolvedExplicitPath"
        }

        if (-not [string]::Equals(
                [IO.Path]::GetFileName($resolvedExplicitPath),
                'TlbImp.exe',
                [StringComparison]::OrdinalIgnoreCase)) {
            throw 'TlbImpPath must identify TlbImp.exe.'
        }

        Assert-Amd64File -Path $resolvedExplicitPath -Description 'TlbImp.exe'
        return $resolvedExplicitPath
    }

    $programFilesX86 = [Environment]::GetFolderPath(
        [Environment+SpecialFolder]::ProgramFilesX86)
    $sdkBin = Join-Path $programFilesX86 'Microsoft SDKs\Windows\v10.0A\bin'
    $candidates = @(
        (Join-Path $sdkBin 'NETFX 4.8.1 Tools\x64\TlbImp.exe'),
        (Join-Path $sdkBin 'NETFX 4.8 Tools\x64\TlbImp.exe'),
        (Join-Path $sdkBin 'NETFX 4.7.2 Tools\x64\TlbImp.exe'),
        (Join-Path $sdkBin 'NETFX 4.6.2 Tools\x64\TlbImp.exe')
    )

    foreach ($candidate in $candidates) {
        if (Test-Path -LiteralPath $candidate -PathType Leaf) {
            Assert-Amd64File -Path $candidate -Description 'TlbImp.exe'
            return [IO.Path]::GetFullPath($candidate)
        }
    }

    $candidateList = $candidates -join [Environment]::NewLine
    throw "An x64 TlbImp.exe was not found in the known Windows SDK locations. Install the .NET Framework SDK or pass -TlbImpPath with an x64 SDK path. Checked:$([Environment]::NewLine)$candidateList"
}

if (-not [Environment]::Is64BitOperatingSystem) {
    throw 'K3D x64 interop generation 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"
    $registeredTypeLibPath = Resolve-RegisteredFilePath `
        -RegistryValue (Get-RegistryValue -BaseKey $machineRoot -SubKey $typeLibKey) `
        -Description 'K3D x64 type library path'
    Assert-Amd64File -Path $registeredTypeLibPath -Description 'K3D type library DLL'

    foreach ($registration in $expectedClasses) {
        Assert-ComClassRegistration `
            -MachineRoot $machineRoot `
            -Registration $registration
    }
}
finally {
    $currentUserRoot.Dispose()
    $machineRoot.Dispose()
}

$resolvedTlbImpPath = Resolve-X64TlbImpPath -ExplicitPath $TlbImpPath
$repositoryRoot = [IO.Path]::GetFullPath((Join-Path $PSScriptRoot '..'))
$outputDirectory = [IO.Path]::GetFullPath(
    (Join-Path $repositoryRoot 'artifacts\K3DInterop'))
$outputPath = Join-Path $outputDirectory 'Interop.K3DAsyncEngineLib.dll'

if (-not $PSCmdlet.ShouldProcess(
        $outputPath,
        'Generate an x64 diagnostic interop assembly from the registered K3D type library')) {
    [pscustomobject][ordered]@{
        Status = 'NotGenerated'
        TypeLibId = $expectedTypeLibId
        TypeLibVersion = $expectedTypeLibVersion
        TypeLibPath = $registeredTypeLibPath
        RegistrationHive = 'HKLM'
        CurrentUserOverrides = 'None'
        ValidatedClasses = @($expectedClasses.Name)
        TlbImpPath = $resolvedTlbImpPath
        OutputPath = $outputPath
        Machine = 'AMD64'
        ComActivated = $false
    }
    return
}

[void] (New-Item -ItemType Directory -Path $outputDirectory -Force)
$stagingDirectory = Join-Path `
    $outputDirectory `
    ('.staging-{0}-{1}' -f $PID, [Guid]::NewGuid().ToString('N'))
$stagedOutputPath = Join-Path $stagingDirectory 'Interop.K3DAsyncEngineLib.dll'
[void] (New-Item -ItemType Directory -Path $stagingDirectory)

try {
    $tlbImpArguments = @(
        $registeredTypeLibPath,
        "/out:$stagedOutputPath",
        '/namespace:K3DAsyncEngineLib',
        '/machine:X64',
        '/silent'
    )

    $tlbImpOutput = & $resolvedTlbImpPath @tlbImpArguments 2>&1
    $tlbImpExitCode = $LASTEXITCODE
    if ($tlbImpExitCode -ne 0) {
        $diagnosticText = ($tlbImpOutput | Out-String).Trim()
        throw "TlbImp.exe failed with exit code $tlbImpExitCode. $diagnosticText"
    }

    if (-not (Test-Path -LiteralPath $stagedOutputPath -PathType Leaf)) {
        throw 'TlbImp.exe reported success but did not create the interop assembly.'
    }

    Assert-Amd64File -Path $stagedOutputPath -Description 'Generated interop assembly'
    [IO.File]::Copy($stagedOutputPath, $outputPath, $true)
    Assert-Amd64File -Path $outputPath -Description 'Final interop assembly'

    $hash = Get-FileHash -LiteralPath $outputPath -Algorithm SHA256
    [pscustomobject][ordered]@{
        Status = 'Generated'
        TypeLibId = $expectedTypeLibId
        TypeLibVersion = $expectedTypeLibVersion
        TypeLibPath = $registeredTypeLibPath
        RegistrationHive = 'HKLM'
        CurrentUserOverrides = 'None'
        ValidatedClasses = @($expectedClasses.Name)
        TlbImpPath = $resolvedTlbImpPath
        OutputPath = $outputPath
        Machine = 'AMD64'
        Sha256 = $hash.Hash
        ComActivated = $false
    }
}
finally {
    if (Test-Path -LiteralPath $stagingDirectory) {
        Remove-Item -LiteralPath $stagingDirectory -Recurse -Force
    }
}