############################
#
# Copyright (c) Microsoft Corporation
#
# Abstract:
#   This script cmdlet enumerates SMAPI MSFT_StorageNodeToPhysicalDisk
#   objects and filters by a StorageNode, PhysicalDisk, or both, if 
#   passed in as parameters.
#
############################

function Get-PhysicalDiskStorageNodeView {
[CmdletBinding()]
param(
    [Microsoft.Management.Infrastructure.CimInstance]
    [PSTypeName("Microsoft.Management.Infrastructure.CimInstance#ROOT/Microsoft/Windows/Storage/MSFT_StorageNode")] 
    [Parameter(
        ValueFromPipeline=$true)]
    $StorageNode = $null,

    [Microsoft.Management.Infrastructure.CimInstance]
    [PSTypeName("Microsoft.Management.Infrastructure.CimInstance#ROOT/Microsoft/Windows/Storage/MSFT_PhysicalDisk")] 
    [Parameter(
        ValueFromPipeline=$true)]
    $PhysicalDisk = $null
    )

## Gather the association objects
$instances = Get-CimInstance -Namespace root\microsoft\windows\storage -ClassName MSFT_StorageNodeToPhysicalDisk

## Filter by StorageNode and PhysicalDisk ObjectIds
ForEach ( $instance in $instances ) {
    if ( $StorageNode ) {
        if ( $instance.StorageNodeObjectId -ne $StorageNode.ObjectId ) {
            continue
        }
    }

    if ( $PhysicalDisk ) {
        if ( $instance.PhysicalDiskObjectId -ne $PhysicalDisk.ObjectId ) {
            continue
        }
    }

    $instance
}
}

function Get-StorageDiagnosticInfo
{
    [CmdletBinding()]
    Param(

        #### -------------------- Parameter sets -------------------------------------####

        [Microsoft.Management.Infrastructure.CimInstance]
        [PSTypeName("Microsoft.Management.Infrastructure.CimInstance#ROOT/Microsoft/Windows/Storage/MSFT_StorageSubSystem")]
        [Parameter(
            ParameterSetName  = "ByStorageSubSystem",
            ValueFromPipeline = $true,
            Mandatory         = $true)]
        [ValidateNotNullOrEmpty()]
        $InputObject,

        [String]
        [Parameter(
            ParameterSetName  = "ByStorageSubSystemFriendlyName",
            Mandatory         = $true,
            Position          = 0)]
        [ValidateNotNullOrEmpty()]
        $StorageSubSystemFriendlyName,

        [String]
        [Parameter(
            ParameterSetName  = "ByStorageSubSystemName",
            Mandatory         = $true)]
        [ValidateNotNullOrEmpty()]
        $StorageSubSystemName,

        [String]
        [Parameter(
            ParameterSetName  = "ByStorageSubSystemUniqueId",
            Mandatory         = $true)]
        [ValidateNotNullOrEmpty()]
        [Alias("StorageSubSystemId")]
        $StorageSubSystemUniqueId,

        #### -------------------- Method parameters ----------------------------------####

        [String]
        [Parameter(
            ParameterSetName = 'ByStorageSubSystem',
            Mandatory        = $true)]
        [Parameter(
            ParameterSetName = 'ByStorageSubSystemFriendlyName',
            Mandatory        = $true)]
        [Parameter(
            ParameterSetName = 'ByStorageSubSystemName',
            Mandatory        = $true)]
        [Parameter(
            ParameterSetName = 'ByStorageSubSystemUniqueId',
            Mandatory        = $true)]
        [ValidateNotNullOrEmpty()]
        [Alias("Path")]
        $DestinationPath,

        [UInt32]
        [Parameter(
            ParameterSetName = 'ByStorageSubSystem',
            Mandatory        = $false)]
        [Parameter(
            ParameterSetName = 'ByStorageSubSystemFriendlyName',
            Mandatory        = $false)]
        [Parameter(
            ParameterSetName = 'ByStorageSubSystemName',
            Mandatory        = $false)]
        [Parameter(
            ParameterSetName = 'ByStorageSubSystemUniqueId',
            Mandatory        = $false)]
        $TimeSpan,

        [String]
        [Parameter(
            ParameterSetName = 'ByStorageSubSystem',
            Mandatory        = $false)]
        [Parameter(
            ParameterSetName = 'ByStorageSubSystemFriendlyName',
            Mandatory        = $false)]
        [Parameter(
            ParameterSetName = 'ByStorageSubSystemName',
            Mandatory        = $false)]
        [Parameter(
            ParameterSetName = 'ByStorageSubSystemUniqueId',
            Mandatory        = $false)]
        $ActivityId,

        [System.Management.Automation.SwitchParameter]
        [Parameter(
            ParameterSetName = 'ByStorageSubSystem',
            Mandatory        = $false)]
        [Parameter(
            ParameterSetName = 'ByStorageSubSystemFriendlyName',
            Mandatory        = $false)]
        [Parameter(
            ParameterSetName = 'ByStorageSubSystemName',
            Mandatory        = $false)]
        [Parameter(
            ParameterSetName = 'ByStorageSubSystemUniqueId',
            Mandatory        = $false)]
        $ExcludeOperationalLog = $false,

        [System.Management.Automation.SwitchParameter]
        [Parameter(
            ParameterSetName = 'ByStorageSubSystem',
            Mandatory        = $false)]
        [Parameter(
            ParameterSetName = 'ByStorageSubSystemFriendlyName',
            Mandatory        = $false)]
        [Parameter(
            ParameterSetName = 'ByStorageSubSystemName',
            Mandatory        = $false)]
        [Parameter(
            ParameterSetName = 'ByStorageSubSystemUniqueId',
            Mandatory        = $false)]
        $ExcludeDiagnosticLog = $false,

        [System.Management.Automation.SwitchParameter]
        [Parameter(
            ParameterSetName = 'ByStorageSubSystem',
            Mandatory        = $false)]
        [Parameter(
            ParameterSetName = 'ByStorageSubSystemFriendlyName',
            Mandatory        = $false)]
        [Parameter(
            ParameterSetName = 'ByStorageSubSystemName',
            Mandatory        = $false)]
        [Parameter(
            ParameterSetName = 'ByStorageSubSystemUniqueId',
            Mandatory        = $false)]
        $IncludeLiveDump = $false,

        #### -------------------- Powershell parameters ------------------------------####

        [Microsoft.Management.Infrastructure.CimSession]
        $CimSession,

        # Provided for compatibility with CDXML cmdlets, not actually used.
        [Int32]
        $ThrottleLimit,

        [Switch]
        $AsJob
    )

    Process
    {
        $info = $resources.info
        $p = $null

        if (-not $CimSession)
        {
            $CimSession = New-CimSession
        }

        if (-not $AsJob)
        {
            Write-Progress -Activity "Get-StorageDiagnosticInfo" -Id 0 -PercentComplete 0 -CurrentOperation "Validating parameters" -Status "0/3"
        }

        ## Gather the instance object

        switch ($psCmdlet.ParameterSetName)
        {
            "ByStorageSubSystem"             { $SubsystemInstance = $InputObject; break; }
            "ByStorageSubSystemFriendlyName" { $SubsystemInstance = Get-StorageSubsystem -FriendlyName $StorageSubSystemFriendlyName -CimSession $CimSession -ErrorAction Stop; break; }
            "ByStorageSubSystemName"         { $SubsystemInstance = Get-StorageSubsystem -Name $StorageSubSystemName -CimSession $CimSession -ErrorAction Stop; break; }
            "ByStorageSubSystemUniqueId"     { $SubsystemInstance = Get-StorageSubsystem -UniqueId $StorageSubSystemUniqueId -CimSession $CimSession -ErrorAction Stop; break; }
        }

        # Would use a closure here, but jobs are run in their own session state.
        $methodblock = {
            param($session, $asjob, $uniqueId, $arguments)

            # Start-Job serializes/deserializes the CimSession,
            # which means it shows up here as having type Deserialized.CimSession.
            # Must recreate or cast the object in order to pass it to Get-CimInstance.
            if ($asjob)
            {
                # We need to load the module or it fails with CMDLET Not Found
                import-module Storage\StorageSubSystem.cdxml

                $session = $session | New-CimSession
            }

            ############################################################################################
            ##
            ## Gather storage diagnostic logs
            ##
            ############################################################################################

            $parameters = @{}

            $parameters.Add("CimSession", $session )
            $parameters.Add("StorageSubSystemUniqueId", $uniqueId )
            $parameters.Add("DestinationPath", $arguments["DestinationPath"] )

            if ($arguments.ContainsKey("TimeSpan"))
            {
                $parameters.Add("TimeSpan", $arguments["TimeSpan"])
            }

            if ($arguments.ContainsKey("ActivityId"))
            {
                $parameters.Add("ActivityId", $arguments["ActivityId"])
            }

            if ($arguments.ContainsKey("ExcludeOperationalLog"))
            {
                $parameters.Add("ExcludeOperationalLog", $arguments["ExcludeOperationalLog"])
            }

            if ($arguments.ContainsKey("ExcludeDiagnosticLog"))
            {
                $parameters.Add("ExcludeDiagnosticLog", $arguments["ExcludeDiagnosticLog"])
            }

            if ($arguments.ContainsKey("IncludeLiveDump"))
            {
                $parameters.Add("IncludeLiveDump", $arguments["IncludeLiveDump"])
            }

            if (-not $asjob)
            {
                Write-Progress -Activity "Get-StorageDiagnosticInfo" -Id 0 -PercentComplete 10 -CurrentOperation "Gathering storage diagnostic logs" -Status "1/3"
            }

            Get-StorageDiagnosticInfoInternal @parameters

            ############################################################################################
            ##
            ## Gather storage objects
            ##
            ############################################################################################

            if (-not $asjob)
            {
                Write-Progress -Activity "Get-StorageDiagnosticInfo" -Id 0 -PercentComplete 60 -CurrentOperation "Gathering storage objects. This may take time depending on the information being gathered." -Status "2/3"
            }

            $globalDatabaseObject = New-Object -TypeName PsObject

            $subsystem = Get-StorageSubsystem -UniqueId $uniqueId -CimSession $CimSession -ErrorAction Stop
            $globalDatabaseObject | Add-Member -NotePropertyName 'Subsystem' -NotePropertyValue $subsystem

            ##
            ## Enumerations via Storage Subsystem
            ##

            $subsystemAssociatedClasses = @{'MSFT_PhysicalDisk'     = 'MSFT_StorageSubSystemToPhysicalDisk'; `
                                            'MSFT_StorageEnclosure' = 'MSFT_StorageSubSystemToStorageEnclosure'; `
                                            'MSFT_StorageNode'      = 'MSFT_StorageSubSystemToStorageNode'; `
                                            'MSFT_StoragePool'      = 'MSFT_StorageSubSystemToStoragePool'; `
                                            'MSFT_VirtualDisk'      = 'MSFT_StorageSubSystemToVirtualDisk'
                                            }

            ForEach ($class in $subsystemAssociatedClasses.Keys)
            {
                if (-not $asjob)
                {
                    Write-Progress -Activity "Gathering storage objects of type $class" -Id 1 -ParentId 0
                }

                $objectTable = @{}

                $objects = Get-CimAssociatedInstance -InputObject $subsystem -Association $subsystemAssociatedClasses[$class] -ResultClassName $class -CimSession $CimSession -ErrorAction Ignore

                ForEach ($object in $objects)
                {
                    $objectTable.Add($object.ObjectId, $object)
                }

                $propertyName = $class.Replace('MSFT_','') + 'Table'
                $globalDatabaseObject | Add-Member -NotePropertyName $propertyName -NotePropertyValue $objectTable
            }

            ##
            ## Storage node views for Disk, Physical Disk and Storage Enclosure
            ##

            $storagenodeViewTypes = @{'PhysicalDisk'     = 'Get-PhysicalDiskStorageNodeView'}

            ForEach ($type in $storagenodeViewTypes.Keys)
            {
                if (-not $asjob)
                {
                    Write-Progress -Activity "Gathering $type storage node view objects" -Id 1 -ParentId 0
                }

                $nodeViewTable = @{}

                $propertyName = $type + 'Table'

                ForEach ($instance in $globalDatabaseObject.$propertyName.Values)
                {
                    $objects = $instance | &$storagenodeViewTypes[$type]
                    $nodeViewTable.Add($instance.ObjectId, $objects)
                }

                $propertyName = $type + 'StorageNodeViewTable'
                $globalDatabaseObject | Add-Member -NotePropertyName $propertyName -NotePropertyValue $nodeViewTable
            }

            ##
            ## Associated storage nodes for Storage Pool, Virtual Disk and Volume
            ##

            $associatedStorageNodeTypes = @{'StoragePool' = 'MSFT_StorageNodeToStoragePool'; `
                                            'VirtualDisk' = 'MSFT_StorageNodeToVirtualDisk'
                                            }

            ForEach ($type in $associatedStorageNodeTypes.Keys)
            {
                if (-not $asjob)
                {
                    Write-Progress -Activity "Gathering $type associated storage node objects" -Id 1 -ParentId 0
                }

                $associatedNodeTable = @{}

                $propertyName = $type + 'Table'

                ForEach ($instance in $globalDatabaseObject.$propertyName.Values)
                {
                    $associatedNodeList = @()

                    $objects = Get-CimAssociatedInstance -InputObject $instance -Association $storagenodeViewTypes[$type] -ResultClassName 'MSFT_StorageNode' -CimSession $CimSession -ErrorAction Ignore

                    ForEach ($object in $objects)
                    {
                        $associatedNodeList += $object.ObjectId
                    }

                    $associatedNodeTable.Add($instance.ObjectId, $associatedNodeList)
                }

                $propertyName = $type + 'AssociatedStorageNodeList'
                $globalDatabaseObject | Add-Member -NotePropertyName $propertyName -NotePropertyValue $associatedNodeTable
            }

            ##
            ## Physical disk associated objects - Storage Fault Domain, Storage Reliability Counter
            ##

            if (-not $asjob)
            {
                Write-Progress -Activity "Gathering Physical Disk associated objects" -Id 1 -ParentId 0
            }

            $physicaldiskTable = @{}

            ForEach ($physicaldisk in $globalDatabaseObject.PhysicalDiskTable.Values)
            {
                $associatedObject = New-Object -TypeName PsObject

                $objects = Get-StorageEnclosure -PhysicalDisk $physicaldisk -CimSession $CimSession

                ForEach ($object in $objects)
                {
                    $type = $object.PSObject.TypeNames[0]
                    $type = $type.SubString($type.LastIndexOf('MSFT_'))
                    $type = $type.Replace('MSFT_','')
                    $associatedObject | Add-Member -NotePropertyName $type -NotePropertyValue $object.ObjectId
                }

                $object = Get-CimAssociatedInstance -InputObject $physicaldisk -Association 'MSFT_PhysicalDiskToStorageReliabilityCounter' -ResultClassName 'MSFT_StorageReliabilityCounter' -CimSession $CimSession -ErrorAction Ignore
                $associatedObject | Add-Member -NotePropertyName 'StorageReliabilityCounter' -NotePropertyValue $object

                $physicaldiskTable.Add($physicaldisk.ObjectId, $associatedObject)
            }

            $globalDatabaseObject | Add-Member -NotePropertyName 'PhysicalDiskAssociationTable' -NotePropertyValue $physicaldiskTable

            ##
            ## Storage pool associated objects - Physical Disk, Virtual Disk, Resiliency Setting, Storage Tier, Storage Job
            ##

            if (-not $asjob)
            {
                Write-Progress -Activity "Gathering Storage Pool associated objects" -Id 1 -ParentId 0
            }

            $storagepoolTable = @{}

            ForEach ($storagepool in $globalDatabaseObject.StoragePoolTable.Values)
            {
                $associatedObject = New-Object -TypeName PsObject

                $physicaldiskList = @()

                $objects = Get-CimAssociatedInstance -InputObject $storagepool -Association 'MSFT_StoragePoolToPhysicalDisk' -ResultClassName 'MSFT_PhysicalDisk' -CimSession $CimSession -ErrorAction Ignore

                ForEach ($object in $objects)
                {
                    $physicaldiskList += $object.ObjectId
                }

                $associatedObject | Add-Member -NotePropertyName 'PhysicalDiskList' -NotePropertyValue $physicaldiskList

                $virtualdiskList = @()

                $objects = Get-CimAssociatedInstance -InputObject $storagepool -Association 'MSFT_StoragePoolToVirtualDisk' -ResultClassName 'MSFT_VirtualDisk' -CimSession $CimSession -ErrorAction Ignore

                ForEach ($object in $objects)
                {
                    $virtualdiskList += $object.ObjectId
                }

                $associatedObject | Add-Member -NotePropertyName 'VirtualDiskList' -NotePropertyValue $virtualdiskList

                $objects = Get-CimAssociatedInstance -InputObject $storagepool -Association 'MSFT_StoragePoolToResiliencySetting' -ResultClassName 'MSFT_ResiliencySetting' -CimSession $CimSession -ErrorAction Ignore
                $associatedObject | Add-Member -NotePropertyName 'ResiliencySetting' -NotePropertyValue $objects

                $objects = Get-CimAssociatedInstance -InputObject $storagepool -Association 'MSFT_StoragePoolToStorageTier' -ResultClassName 'MSFT_StorageTier' -CimSession $CimSession -ErrorAction Ignore
                $associatedObject | Add-Member -NotePropertyName 'StorageTier' -NotePropertyValue $objects

                $objects = Get-CimAssociatedInstance -InputObject $storagepool -Association 'MSFT_StorageJobToAffectedStorageObject' -ResultClassName 'MSFT_StorageJob' -CimSession $CimSession -ErrorAction Ignore
                $associatedObject | Add-Member -NotePropertyName 'StorageJob' -NotePropertyValue $objects

                $storagepoolTable.Add($storagepool.ObjectId, $associatedObject)
            }

            $globalDatabaseObject | Add-Member -NotePropertyName 'StoragePoolAssociationTable' -NotePropertyValue $storagepoolTable

            ##
            ## Virtual disk associated objects - Physical Disk, Disk, Storage Tier, Storage Job
            ##

            if (-not $asjob)
            {
                Write-Progress -Activity "Gathering Virtual Disk associated objects" -Id 1 -ParentId 0
            }

            $virtualdiskTable = @{}

            ForEach ($virtualdisk in $globalDatabaseObject.VirtualDiskTable.Values)
            {
                $associatedObject = New-Object -TypeName PsObject

                $physicaldiskList = @()

                $objects = Get-CimAssociatedInstance -InputObject $virtualdisk -Association 'MSFT_VirtualDiskToPhysicalDisk' -ResultClassName 'MSFT_PhysicalDisk' -CimSession $CimSession -ErrorAction Ignore

                ForEach ($object in $objects)
                {
                    $physicaldiskList += $object.ObjectId
                }

                $associatedObject | Add-Member -NotePropertyName 'PhysicalDiskList' -NotePropertyValue $physicaldiskList

                $object = Get-CimAssociatedInstance -InputObject $virtualdisk -Association 'MSFT_VirtualDiskToDisk' -ResultClassName 'MSFT_Disk' -CimSession $CimSession -ErrorAction Ignore
                $associatedObject | Add-Member -NotePropertyName 'Disk' -NotePropertyValue $object.ObjectId

                $objects = Get-CimAssociatedInstance -InputObject $virtualdisk -Association 'MSFT_VirtualDiskToStorageTier' -ResultClassName 'MSFT_StorageTier' -CimSession $CimSession -ErrorAction Ignore
                $associatedObject | Add-Member -NotePropertyName 'StorageTier' -NotePropertyValue $objects

                $objects = Get-CimAssociatedInstance -InputObject $virtualdisk -Association 'MSFT_StorageJobToAffectedStorageObject' -ResultClassName 'MSFT_StorageJob' -CimSession $CimSession -ErrorAction Ignore
                $associatedObject | Add-Member -NotePropertyName 'StorageJob' -NotePropertyValue $objects

                $virtualdiskTable.Add($virtualdisk.ObjectId, $associatedObject)
            }

            $globalDatabaseObject | Add-Member -NotePropertyName 'VirtualDiskAssociationTable' -NotePropertyValue $virtualdiskTable

            ##
            ## Write to xml database file
            ##

            if (-not $asjob)
            {
                Write-Progress -Activity "Get-StorageDiagnosticInfo" -Id 0 -PercentComplete 95 -CurrentOperation "Saving storage objects" -Status "3/3"
            }

            $databaseFilePath = $arguments["DestinationPath"] +'\StorageObjects.xml'
            $globalDatabaseObject | Export-Clixml -Depth 999 -Path $databaseFilePath

            if (-not $asjob)
            {
                Write-Progress -Activity "Get-StorageDiagnosticInfo" -Completed -Status "3/3"
            }
        }

        if ($asjob)
        {
            $p = $true

            Start-Job -Name GetDiagnosticInfo `
                      -ScriptBlock $methodblock `
                      -ArgumentList @($CimSession, $p, $SubsystemInstance.UniqueId, $PSBoundParameters)
        }
        else
        {
            &$methodblock $CimSession $p $SubsystemInstance.UniqueId $PSBoundParameters
        }
    }
}
New-Alias Get-PhysicalDiskSNV Get-PhysicalDiskStorageNodeView

Export-ModuleMember -Alias * -Function *