Import-Project.ps1

# Copyright (c) Matthias Wolf, Mawosoft.

using namespace System
using namespace System.Collections.Generic
using namespace System.IO
using namespace System.Reflection
using namespace System.Management.Automation

<#
.SYNOPSIS
    Imports a collection of PowerShell script files as a project.
.DESCRIPTION
    TODO details about source file handling and module files psd1/psm1.
    TODO explain type import via accels.
.INPUTS
    You can pipe paths to this cmdlet.
.OUTPUTS
    None - By default, this cmdlet returns no output.
    [ExtendedModuleInfo] - If PassThru has been specified.
#>

function Import-Project {
    [CmdletBinding(PositionalBinding = $false, DefaultParameterSetName = 'Path')]
    param(
        [Alias('LP', 'PSPath')]
        [Parameter(Mandatory, ParameterSetName = 'LiteralPath', ValueFromPipelineByPropertyName, Position = 0)]
        [ValidateNotNullOrEmpty()]
        # Specifies paths to PowerShell script files. Wildcards are not permitted.
        [string[]]$LiteralPath,

        [Parameter(Mandatory, ParameterSetName = 'Path', ValueFromPipeline, Position = 0)]
        [ValidateNotNullOrEmpty()]
        [SupportsWildcards()]
        # Specifies paths to PowerShell script files that may contain wildcards.
        [string[]]$Path,

        # Specifies the name for a script project. The default is 'ScriptProject'.
        # If the provided paths contain a module file, the module name is used instead.
        [string]$Name,

        [ValidateSet('Global', 'TypeAccelerator')]
        # Specifies how types (PowerShell classes) are imported from the project. Accepted values are:
        # 'Global' - Types are imported into the global scope.
        # 'TypeAccelerator' - Types are imported as user-defined type accelerators.
        [string]$TypeImport = 'TypeAccelerator',

        [ValidateSet('All', 'Script')]
        # Specifies how variables are imported from the project. Accepted values are:
        # 'All' - All non-local variables are imported.
        # 'Script' - Only variables explicitly prefixed with 'script:' are imported.
        [string]$VariableImport = 'Script',

        [ValidateSet('All', 'Open', 'Marked', 'None')]
        # Specifies how PSES diagnostics for script files are updated. Accepted values are:
        # 'All' - Updates all files belonging to the project.
        # 'Open' - Updates only files that are opened in the editor.
        # 'Marked' - Updates only files that already contain diagnostic markers.
        # 'None' - No diagnostics will be updated.
        [string]$UpdateDiagnostic = 'Marked',

        # Specifies that the cmdlet should return the [ExtendedModuleInfo] object for the project.
        [switch]$PassThru,

        # Specifies that the cmdlet should proceed even if PSES is not available.
        [switch]$Force,

        # If the project is a PowerShell module, specifies optional arguments,
        # passed via the Import-Module cmdlet.
        [Alias('Args')]
        [Object[]]$ArgumentList
    )
    begin {
        if ($script:PSES.Error) {
            if ($Force) {
                Write-Warning $script:PSES.Error
            }
            else {
                $PSCmdlet.ThrowTerminatingError([ErrorRecord]::new(
                        [InvalidOperationException]::new($script:PSES.Error),
                        $script:PSES.Error, [ErrorCategory]::InvalidOperation, $null))
                return
            }
        }
        $pathIntrinsics = $ExecutionContext.SessionState.Path
        $currentDir = $pathIntrinsics.CurrentFileSystemLocation
        $resolvedPaths = [List[string]]::new()
        $sourceFiles = [List[string]]::new()
        $seen = [HashSet[string]]::new($script:PathComparer)
        $psm1Files = [List[string]]::new()
        $modulePath = ''
    }
    process {
        $resolvedPaths.Clear()
        if ($LiteralPath) {
            foreach ($p in $LiteralPath) {
                $p2 = [Path]::GetFullPath([Path]::Combine($currentDir, $p))
                if ($seen.Add($p2)) { $resolvedPaths.Add($p2) }
            }
        }
        if ($Path) {
            foreach ($p in $Path) {
                foreach ($p2 in $pathIntrinsics.GetResolvedPSPathFromPSPath($p)) {
                    if ($seen.Add($p2.Path)) { $resolvedPaths.Add($p2.Path) }
                }
            }
        }
        foreach ($p in $resolvedPaths) {
            $ext = [Path]::GetExtension($p)
            if (-not [File]::Exists($p)) {
                Write-Error "File not found: $p"
            }
            elseif ($ext -eq '.psd1') {
                if ($modulePath) {
                    Write-Error "Multiple manifest files. Ignoring: $p"
                }
                else {
                    $modulePath = $p
                }
            }
            elseif ($ext -eq '.psm1') {
                $psm1Files.Add($p)
            }
            elseif ($ext -eq '.ps1') {
                $sourceFiles.Add($p)
            }
            else {
                Write-Error "Invalid file extension: $p"
            }
        }
    }
    end {
        if ($modulePath) {
            $sourceFiles.AddRange($psm1Files)
        }
        elseif ($psm1Files) {
            $modulePath = $psm1Files[0]
            $sourceFiles.AddRange($psm1Files)
            for ($i = 1; $i -lt $psm1Files.Count; $i++) {
                Write-Error "Multiple module files. Not importing: $($psm1Files[$i])"
            }
        }
        else {
            if ($sourceFiles.Count -eq 0) {
                throw 'No source files found.'
                return
            }
            if ($ArgumentList) {
                Write-Error "The parameter 'ArgumentList' is only applicable to modules."
            }
        }
        if ($modulePath) {
            $Name = [Path]::GetFileNameWithoutExtension($modulePath)
        }
        elseif (-not $Name) {
            $Name = 'ScriptProject'
        }
        $reflector = $script:Reflector
        $SMA = $script:SMA
        $context = $reflector.Expose($ExecutionContext)._context

        $module = Get-Module -Name $Name
        if ($module) {
            Remove-Project -ModuleInfo $module
        }
        if ($modulePath) {
            $module = Import-Module -Name $modulePath -ArgumentList $ArgumentList -Global -PassThru
        }
        else {
            # New-Module has no -Global switch. Any exports would be imported in our own module scope.
            # While we can prevent default exports by adding Export-ModuleMember, we cannot prevent
            # explicit exports from dot-sourced scripts.
            # We are using New-Module here instead of [psmoduleinfo]::new() to easily define a name.
            $module = New-Module -Name $Name -ScriptBlock {}
            $sb = {
                if ($args.Length -gt 1 -and $null -ne $args[1]) {
                    $ErrorActionPreference = $args[1]
                }
                foreach ($arg in $args[0]) { . $arg }
                Remove-Variable 'arg', 'ErrorActionPreference' -ErrorAction Ignore
                Export-ModuleMember # No default exports
            }
            $sb = $module.NewBoundScriptBlock($sb)
            $propagate = $context.PropagateExceptionsToEnclosingStatementBlock
            try {
                # We want non-terminating errors not to throw if the EAP allows it.
                $context.PropagateExceptionsToEnclosingStatementBlock = $false
                . $sb $sourceFiles $PSBoundParameters['ErrorAction'] # Dot-sourced into $module
            }
            catch {
                # We only need to restore here before re-throwing.
                # Auto-restore happens internally after the try-catch block.
                $context.PropagateExceptionsToEnclosingStatementBlock = $propagate
                throw
                return
            }
            Write-Verbose "Creating module: $Name"
            $module = Import-Module -ModuleInfo $module -Global -PassThru
        }
        if (-not $module) {
            # Someone better has reported an error.
            return
        }
        $moduleInfo = Get-ExtendedModuleInfo -ModuleInfo $module -Merge -Parse:($VariableImport -eq 'Script')
        foreach ($p in $moduleInfo.ScriptFiles) {
            if ($seen.Add($p)) { $sourceFiles.Add($p) }
        }
        $topLevelSsi = $context.TopLevelSessionState
        $globalScope = $topLevelSsi.GlobalScope
        if ($TypeImport -eq 'TypeAccelerator') {
            foreach ($kvp in $moduleInfo.Types.GetEnumerator()) {
                $SMA.TypeAccelerators::Add($kvp.Key, $kvp.Value)
                Write-Verbose "Importing type accelerator: $($kvp.Key)"
            }
        }
        else {
            if ($moduleInfo.Types.Count -ne 0) {
                foreach ($kvp in $moduleInfo.Types.GetEnumerator()) {
                    $globalScope.AddType($kvp.Key, $kvp.Value)
                    Write-Verbose "Importing global type: $($kvp.Key)"
                }
                $trs = $globalScope.TypeResolutionState
                $assemblies = [HashSet[Assembly]]::new($trs.assemblies)
                $assemblies.UnionWith($moduleInfo.Assemblies)
                $trs = $SMA.TypeResolutionState.new($trs.namespaces, $assemblies)
                $globalScope.TypeResolutionState = $trs.CloneWithAddTypesDefined($globalScope.typeTable.Keys)
                Write-Verbose 'Updating global type resolution state.'
            }
        }
        foreach ($kvp in $moduleInfo.Functions.GetEnumerator()) {
            $f = $kvp.Value
            $f = $globalScope.SetFunction($kvp.Key, $f.ScriptBlock, $f, $f.Options, <# force: #> $true, [CommandOrigin]::Internal, $topLevelSsi.ExecutionContext)
            Write-Verbose "Importing global function: $($kvp.Key)"
        }
        $module = $moduleInfo.Module
        if ($moduleInfo.Variables.Count -ne 0 -and $null -ne $module.SessionState) {
            $moduleSsi = $reflector.Expose($module.SessionState).Internal
            [bool]$allVariables = $VariableImport -eq 'All'
            foreach ($kvp in $moduleInfo.Variables.GetEnumerator()) {
                if ($allVariables -or $moduleInfo.ScriptVariables.Contains($kvp.Key)) {
                    $v = $kvp.Value
                    $reflector.Expose($v).SetModule($module)
                    $moduleSsi.ExportedVariables.Add($v)
                    $v = $globalScope.NewVariable($v, <# force: #> $true, $moduleSsi)
                    Write-Verbose "Importing global variable: $($kvp.Key)"
                }
            }
            $exported = $moduleSsi.ExportedVariables
            $exported.Sort({ param($x, $y) return [string]::Compare($x.Name, $y.Name, [StringComparison]::OrdinalIgnoreCase) })
            for ($i = $exported.Count - 2; $i -ge 0; $i--) {
                if ([string]::Equals($exported[$i].Name, $exported[$i + 1].Name, [StringComparison]::OrdinalIgnoreCase)) {
                    $exported.RemoveAt($i + 1)
                }
            }
        }
        if ($UpdateDiagnostic -ne 'None' -and -not $script:PSES.Error) {
            Update-EditorDiagnostic -ResolvedPath $sourceFiles -UpdateDiagnostic $UpdateDiagnostic
        }
        if ($PassThru) {
            $moduleInfo
        }
    }
}