$ReadActions = @(
    "*/read"
)

$ReadWriteActions = @(
    "*/read",
    "Microsoft.Compute/virtualMachines/start/action",
    "Microsoft.Compute/virtualMachines/deallocate/action",
    "Microsoft.Compute/virtualMachineScaleSets/start/action",
    "Microsoft.Compute/virtualMachineScaleSets/deallocate/action",
    "Microsoft.Compute/virtualMachineScaleSets/scale/action",
    "Microsoft.Compute/virtualMachineScaleSets/write",
    "Microsoft.Compute/virtualMachineScaleSets/delete",
    "Microsoft.Compute/virtualMachines/delete",
    "Microsoft.Compute/virtualMachines/write",
    "Microsoft.DBforPostgreSQL/flexibleServers/start/action",
    "Microsoft.DBforPostgreSQL/flexibleServers/stop/action",
    "Microsoft.DBforMySQL/flexibleServers/start/action",
    "Microsoft.DBforMySQL/flexibleServers/stop/action",
    "Microsoft.DBforMySQL/servers/start/action",
    "Microsoft.DBforMySQL/servers/stop/action",
    "Microsoft.DBforMariaDB/servers/start/action",
    "Microsoft.DBforMariaDB/servers/stop/action",
    "Microsoft.Compute/disks/delete",
    "Microsoft.Compute/snapshots/delete",
    "Microsoft.Compute/images/delete",
    "Microsoft.Network/publicIPAddresses/delete",
    "Microsoft.Resources/tags/read",
    "Microsoft.Resources/tags/write"
)


function Setup-UniskaiConnection {
    param (
        [Parameter(Mandatory = $false)]
        [string]$SubscriptionName,

        [Parameter(Mandatory = $false)]
        [string]$SubscriptionId,

        [Parameter(Mandatory = $false)]
        [string]$ApplicationName,

        [Parameter(Mandatory = $false)]
        [string]$CustomRoleName,

        [Parameter(Mandatory = $false)]
        [string]$CustomRoleAccess,
    
        [Parameter(Mandatory = $false)]
        [string]$PrincipalId,
    
        [Parameter(Mandatory = $false)]
        [bool]$CreateSpCredential = $true
    )

    # Check if neither or both SubscriptionName and SubscriptionId are provided
    if ((-not $SubscriptionName -and -not $SubscriptionId) -or ($SubscriptionName -and $SubscriptionId)) {
        Write-Error "Please specify either SubscriptionName or SubscriptionId, but not both."
        return
    }

    # Wrap the whole script in a try-catch block to catch exceptions
    try {
        $subscriptions = Get-AzSubscription
        if (-not $subscriptions) {
            Write-Error "No subscriptions found. Please check your Azure account."
            return
        }

        # Choose your subscription by using 'SubscriptionName' or 'SubscriptionId'
        if ($SubscriptionName) {
            $subscription = $subscriptions | Where-Object {$_.Name -eq $SubscriptionName} | Select-Object -First 1
        } elseif ($SubscriptionId) {
            $subscription = $subscriptions | Where-Object {$_.Id -eq $SubscriptionId} | Select-Object -First 1
        }
        if (-not $subscription) {
            Write-Error "Subscription '$SubscriptionName' or '$SubscriptionId' not found. Please provide your subscription name or id."
            return
        }
        Set-AzContext -Subscription $subscription

        $tenantId = $subscription.TenantId
        $uniqueSuffix = (Get-Date).ToString("yyyyMMddHHmm")

        # If the app client name is not passed, use the default name
        if (-not $ApplicationName) {
            $ApplicationName = "Uniskai"
            Write-Output "Using name '$ApplicationName' as no ApplicationName parameter provided."
        }

        # Check if the application already exists
        if ($PrincipalId) {
            $principal = Get-AzADServicePrincipal -ObjectId $PrincipalId | Select-Object -First 1
            if (-not $principal) {
                Write-Error "Service principal with id $PrincipalId not found. Please check your service principal id."
                return
            }
        } else {
            $principal = Get-AzADServicePrincipal -DisplayName $ApplicationName | Select-Object -First 1
            if (-not $principal) {
                Write-Output "Creating new service principal $ApplicationName"
                $principal = New-AzADServicePrincipal -DisplayName $ApplicationName
                if (-not $principal) {
                    Write-Error "Failed to create service principal. Please check your permissions to create AD service principals and try again."
                    Write-Warning "You can create the service principal manually and pass the PrincipalId parameter to this script: -PrincipalId <id>."
                    return
                }

                Write-Output "Waiting 20 seconds for the service principal to propagate in other Azure services..."
                Start-Sleep -Seconds 20
            }
        }

        $connectionConfig = @{
            tenant_id = $tenantId
            subscription_id = $subscription.Id
            principal_id = $principal.Id
            client_id = $principal.AppId
        }

        if ($CreateSpCredential) {
            Write-Output "Creating new secret for service principal $ApplicationName ($($principal.Id)) in subscription $subscription."
            $secret = New-AzADSpCredential -ObjectId $principal.Id -StartDate (Get-Date) -EndDate (Get-Date).AddMonths(24)
            if (-not $secret) {
                Write-Error "Failed to create a new secret. Please check your permissions to create credentials for principal $ApplicationName with ID $($principal.Id)."
                Write-Warning "Consider customizing the application name by passing the ApplicationName parameter to this script: -ApplicationName <name>."
                return
            }
            Write-Output "Created new secret for service principal $ApplicationName ($($principal.Id)) in subscription $subscription with KeyId $($secret.KeyId)."
            $connectionConfig["client_secret"] = $secret.SecretText
        }
 
        Write-Output "Connection configuration:"
        Write-Output $($connectionConfig | ConvertTo-Json)

        $roleName = "Reader"
        $customRoleActions = $null

        if ($CustomRoleAccess) {
            if (-not $CustomRoleName) {
                $roleName = (@($ApplicationName, $subscription.Id, $CustomRoleAccess) | Where-Object { $_ }) -join '-'
            }
            if ($CustomRoleAccess -eq "Read") {
                $customRoleActions = $ReadActions
            } elseif ($CustomRoleAccess -eq "ReadWrite") {
                $customRoleActions = $ReadWriteActions
            } else {
                Write-Error "Invalid value for CustomRoleAccess. Please use either 'Read' or 'ReadWrite'."
                return
            }
        }

        # Check if the custom role is already created
        if ($customRoleActions) {
            $roleDefinition = Get-AzRoleDefinition -Name $roleName
            if (-not $roleDefinition) {
                Write-Output "Creating new custom role $roleName in subscription $subscription."

                $roleDefinition = Get-AzRoleDefinition -Name "Reader"
                $roleDefinition.Id = $null
                $roleDefinition.Name = $roleName
                $roleDefinition.Description = "$CustomRoleAccess role for Uniskai."
                $roleDefinition.IsCustom = $True
                $roleDefinition.Actions = $customRoleActions
                $roleDefinition.AssignableScopes = @("/subscriptions/$($subscription.Id)")
                $roleDefinition = New-AzRoleDefinition -Role $roleDefinition
                if (-not $roleDefinition) {
                    Write-Error "Failed to create custom role $roleName."
                    Write-Error "Please check your permissions to create custom roles within subscription $subscription and try again."
                    return
                }
                $roleName = $roleDefinition.Name
                Write-Output "Waiting 10 seconds for the role definition to propagate in Azure services..."
                Start-Sleep -Seconds 10
            } elseif (Compare-Object -ReferenceObject $roleDefinition.Actions -DifferenceObject $customRoleActions) {
                $roleDefinition.Actions = $customRoleActions
                $roleDefinition.Description = "$CustomRoleAccess role for Uniskai."
                $roleDefinition.AssignableScopes.Add("/subscriptions/$($subscription.Id)")
                Write-Output "Updating custom role $roleName in subscription $subscription."
                $roleDefinition = Set-AzRoleDefinition -Role $roleDefinition
            }
        }

        # Check if the role is already assigned
        $roleAssignment = Get-AzRoleAssignment -ObjectId $principal.Id -Scope "/subscriptions/$($subscription.Id)" `
            | Where-Object {$_.RoleDefinitionName -eq $roleName} | Select-Object -First 1
        if (-not $roleAssignment) {
            # If not, assign the role to the new client
            Write-Output "Assigning $roleName role to $ApplicationName application in $subscription."
            $roleAssignment = New-AzRoleAssignment -ObjectId $principal.Id -RoleDefinitionName $roleName -Scope "/subscriptions/$($subscription.Id)"
            if (-not $roleAssignment) {
                Write-Error "Failed to assign $roleName role. Please check your permissions to assign roles within subscription $subscription and try again."
                return
            }
        }

        if ($CreateSpCredential) {
            $configFile = "uniskai-connection-$uniqueSuffix.json"
            $connectionConfig | ConvertTo-Json | Out-File -FilePath $configFile
        
            Write-Output "Connection configuration file created: $configFile"
            Write-Output "Download this file and register it in Uniskai."
            Write-Warning "If you are using Azure PowerShell, click on the Upload/Download files button at the toolbar,"
            Write-Warning "enter the '$configFile' file name and click on Download."

            Read-Host "After registering the file in Uniskai, press any key to continue..."
            # Remove the connection configuration file
            # Write-Output "Removing connection configuration file..."
            # Remove-Item -Path $configFile -Force
        }
    } catch {
        Write-Error $_.Exception.Message
    }
}