#Requires -Version 7.0
#Requires -Module jwtPS
function Connect-AppleBusinessManager {
    [CmdletBinding(DefaultParameterSetName = 'EnvironmentVariable')]
    param(
        [string][Parameter(ParameterSetName = 'PrivateKeyAsString', Mandatory)]$ClientId,
        [string][Parameter(ParameterSetName = 'PrivateKeyAsString', Mandatory)]$PrivateKey,
        [string][Parameter(ParameterSetName = 'PrivateKeyAsString', Mandatory)]$PrivateKeyId
    )

    if ($PSBoundParameters.Count -eq 0) {
        # Parameters were not bound, check to see if they're already set before attempting to authenticate
        Write-Verbose "Checking to see if existing credentials were set at a script scope that can be re-used"
        if (-not $Script:ClientId -or -not $Script:PrivateKey -or -not $Script:PrivateKeyId) {
            Write-Verbose "Existing credentials not loaded, checking for environment variables"
            if ($PSCmdlet.ParameterSetName -eq 'EnvironmentVariable') {
                if (-not $Env:AppleBusinessManagerClientId -or -not $Env:AppleBusinessManagerPrivateKeyId -or -not $Env:AppleBusinessManagerPrivateKey) {
                    throw "Client ID, Private Key ID and Private Key environment variables were not set for Apple Business Manager"
                }
                $Script:ClientId = $Env:AppleBusinessManagerClientId
                $Script:PrivateKey = $Env:AppleBusinessManagerPrivateKey
                $Script:PrivateKeyId = $Env:AppleBusinessManagerPrivateKeyId
            }
        }
    }
    
    if ($PSCmdlet.ParameterSetName -eq 'PrivateKeyAsString') {
        # Assign the parameters to script scope so we can reuse them later
        $Script:ClientId = $ClientId
        $Script:PrivateKey = $PrivateKey
        $Script:PrivateKeyId = $PrivateKeyId
    }

    $Header = @{
        'kid' = $Script:PrivateKeyId
    }

    $Payload = @{
        aud = "https://account.apple.com/auth/oauth2/v2/token" 
        iss = $Script:ClientId
        sub = $Script:ClientId
        iat = ([System.DateTimeOffset]::Now).ToUnixTimeSeconds()
        exp = ([System.DateTimeOffset]::Now.AddMinutes(15)).ToUnixTimeSeconds()
        jti = [guid]::NewGuid()
    }

    $Hashing = [jwtTypes+encryption]::SHA256
    $Signature = [jwtTypes+algorithm]::ECDsa
    $Algorithm = [jwtTypes+cryptographyType]::new($Signature, $Hashing)
    $JWT = New-JWT -Payload $Payload -Algorithm $Algorithm -Secret $Script:PrivateKey -Header $Header
    Write-Verbose $JWT

    $Script:Body = @{
        'grant_type'            = 'client_credentials'
        'client_id'             = $Script:ClientId
        'client_assertion'      = $JWT
        'client_assertion_type' = 'urn:ietf:params:oauth:client-assertion-type:jwt-bearer'
        'scope'                 = 'business.api'
    }

    $Script:AuthResponse = Invoke-RestMethod -Method Post -Uri 'https://account.apple.com/auth/oauth2/token' -Body $Script:Body -SkipHttpErrorCheck
    if ($Script:AuthResponse.error) {
        throw $Script:AuthResponse.Error
    }
    $Script:ExpiresAt = (Get-Date).AddSeconds($Script:AuthResponse.expires_in)
}

function Get-AppleBusinessManagerBearerToken {
    if (-not $Script:AuthResponse) {
        try {
            Connect-AppleBusinessManager
        }
        catch {
            throw "Authorization has not been completed, use Connect-AppleBusinessManager first."
        }
    }
    if ((Get-Date).AddMinutes(1) -ge $Script:ExpiresAt) {
        # Access token is approaching expiration, get a new access token
        Connect-AppleBusinessManager
    }

    return ($Script:AuthResponse.access_token | ConvertTo-SecureString -AsPlainText -Force)
}

function Invoke-AppleBusinessManagerPagedApiRequest {
    param (
        [Parameter(Mandatory = $true)][uri] $Uri
    )
    $Results = New-Object System.Collections.ArrayList
    while ($Uri) {
        Write-Verbose "Making request to $Uri"
        $Result = Invoke-RestMethod $Uri -Authentication Bearer -Token (Get-AppleBusinessManagerBearerToken) -ErrorAction Stop
        $Uri = $Result.links.next
        $Results.AddRange(@($Result.data)) | Out-Null
    }
    return $Results
}

function Get-AppleBusinessManagerOrgDevice {
    [Alias('Get-AppleBusinessManagerOrgDevices')]
    param (
        [Parameter(Mandatory, ParameterSetName = 'Read')][string] $OrgDeviceId,
        [Parameter(ParameterSetName = 'Read')]
        [Parameter(ParameterSetName = 'List')]
        [string[]] $Fields,
        [Parameter( ParameterSetName = 'Read')]
        [Parameter(ParameterSetName = 'List')]
        [int] $Limit
    )
    $Uri = switch ($PSCmdlet.ParameterSetName) {
        'Read' { "https://api-business.apple.com/v1/orgDevices/$([System.Web.HttpUtility]::UrlEncode($OrgDeviceId))" }
        Default { "https://api-business.apple.com/v1/orgDevices" }
    }
    $UriBuilder = [System.UriBuilder]::new($Uri)
    $QueryString = [System.Web.HttpUtility]::ParseQueryString($UriBuilder.Query)
    if ($PSBoundParameters.ContainsKey('Fields')) {
        $QueryString.Set('fields[orgDevices]', $Fields -join ',')
    }
    if ($PSBoundParameters.ContainsKey('Limit')) {
        $QueryString.Set('limit', $Limit)
    }
    $UriBuilder.Query = $QueryString.ToString()
    return Invoke-AppleBusinessManagerPagedApiRequest -Uri $UriBuilder.Uri
}

function Get-AppleBusinessManagerOrgDeviceMdmServerId {
    param (
        [Parameter(Mandatory)]
        [string] $OrgDeviceId
    )
    return Invoke-AppleBusinessManagerPagedApiRequest -Uri "https://api-business.apple.com/v1/orgDevices/$([System.Web.HttpUtility]::UrlEncode($OrgDeviceId))/relationships/assignedServer"
}

function Get-AppleBusinessManagerOrgDeviceMdmServer {
    param (
        [Parameter(Mandatory)]
        [string] $OrgDeviceId,
        [string[]] $Fields
    )
    $UriBuilder = [System.UriBuilder]::new("https://api-business.apple.com/v1/orgDevices/$([System.Web.HttpUtility]::UrlEncode($OrgDeviceId))/assignedServer")
    $QueryString = [System.Web.HttpUtility]::ParseQueryString($UriBuilder.Query)
    if ($PSBoundParameters.ContainsKey('Fields')) {
        $QueryString.Set('fields[mdmServers]', $Fields -join ',')
    }
    $UriBuilder.Query = $QueryString.ToString()
    return Invoke-AppleBusinessManagerPagedApiRequest -Uri $UriBuilder.Uri
}

function Get-AppleBusinessManagerMdmServer {
    [Alias('Get-AppleBusinessManagerMdmServers')]
    param (
        [string[]] $Fields,
        [int] $Limit
    )
    $UriBuilder = [System.UriBuilder]::new("https://api-business.apple.com/v1/mdmServers")
    $QueryString = [System.Web.HttpUtility]::ParseQueryString($UriBuilder.Query)
    if ($PSBoundParameters.ContainsKey('Fields')) {
        $QueryString.Set('fields[mdmServers]', $Fields -join ',')
    }
    if ($PSBoundParameters.ContainsKey('Limit')) {
        $QueryString.Set('limit', $Limit)
    }
    $UriBuilder.Query = $QueryString.ToString()
    return Invoke-AppleBusinessManagerPagedApiRequest -Uri $UriBuilder.Uri
}

function Get-AppleBusinessManagerMdmServerDevice {
    param (
        [Parameter(Mandatory)]
        [string] $MdmServerId,
        [int] $Limit
    )
    $UriBuilder = [System.UriBuilder]::new("https://api-business.apple.com/v1/mdmServers/$([System.Web.HttpUtility]::UrlEncode($MdmServerId))/relationships/devices")
    $QueryString = [System.Web.HttpUtility]::ParseQueryString($UriBuilder.Query)
    if ($PSBoundParameters.ContainsKey('Limit')) {
        $QueryString.Set('limit', $Limit)
    }
    $UriBuilder.Query = $QueryString.ToString()
    return Invoke-AppleBusinessManagerPagedApiRequest -Uri $UriBuilder.Uri
}