%pip install semantic-link-labs

import sempy_labs as labs
import sempy_labs.admin as admin

azure_subscription_id = '' # Enter your Azure subscription ID
resource_group = '' # Enter the name of the resource group (to be used to create the new F skus)

key_vault_uri = '', # Enter your key vault URI
key_vault_tenant_id = '', # Enter the key vault key to the secret storing your Tenant ID
key_vault_client_id = '', # Enter the key vault key to the secret storing your Client ID (Application ID)
key_vault_client_secret = '' # Enter the key vault key to the secret storing your Client Secret

with labs.service_principal_authentication(
    key_vault_uri=key_vault_uri, 
    key_vault_tenant_id=key_vault_tenant_id,
    key_vault_client_id=key_vault_client_id,
    key_vault_client_secret=key_vault_client_secret):

    labs.migrate_capacities(
        azure_subscription_id = azure_subscription_id,
        resource_group = resource_group,
        capacities = 'CapacityA',
        p_sku_only = True,
    )

with labs.service_principal_authentication(
    key_vault_uri=key_vault_uri, 
    key_vault_tenant_id=key_vault_tenant_id,
    key_vault_client_id=key_vault_client_id,
    key_vault_client_secret=key_vault_client_secret):

    labs.migrate_capacities(
        azure_subscription_id = azure_subscription_id,
        resource_group = resource_group,
        capacities = ['CapacityA', 'CapacityB', 'CapacityC'],
        p_sku_only = True,
    )

with labs.service_principal_authentication(
    key_vault_uri=key_vault_uri, 
    key_vault_tenant_id=key_vault_tenant_id,
    key_vault_client_id=key_vault_client_id,
    key_vault_client_secret=key_vault_client_secret):

    labs.migrate_capacities(
        azure_subscription_id = azure_subscription_id,
        resource_group = resource_group,
        capacities = None,
        p_sku_only = True,
    )

resource_group_mapping = {
    "CapacityA": "ResourceGroupA",
    "CapacityB": "ResourceGroupA",
    "CapacityC": "ResourceGroupB",
}

with labs.service_principal_authentication(
    key_vault_uri=key_vault_uri, 
    key_vault_tenant_id=key_vault_tenant_id,
    key_vault_client_id=key_vault_client_id,
    key_vault_client_secret=key_vault_client_secret):

    labs.migrate_capacities(
        azure_subscription_id = azure_subscription_id,
        resource_group = resource_group_mapping,
        capacities = ['CapacityA', 'CapacityB', 'CapacityC'],
        p_sku_only = True,
    )

source_capacity = '' # Enter the Source capacity name
target_capacity = '' # Enter the Target capacity name (already exists) 

admin.assign_workspaces_to_capacity(
    source_capacity=source_capacity,
    target_capacity=target_capacity,
    workspace=None,
)

capacity_mapping = {
    "capacitya": "capacityanew", # Format is "Source SKU": "Target SKU"
    "capacityb": "capacitybnew",
    "capacityc": "capacitycnew",
}

capacities = list(capacity_mapping.keys())

for capacity in capacities:
    admin.assign_workspaces_to_capacity(
        source_capacity=capacity,
        target_capacity=capacity_mapping.get(capacity),
        workspace=None,
    )

with labs.service_principal_authentication(
    key_vault_uri=key_vault_uri, 
    key_vault_tenant_id=key_vault_tenant_id,
    key_vault_client_id=key_vault_client_id,
    key_vault_client_secret=key_vault_client_secret):

    labs.migrate_fabric_trial_capacity(
        azure_subscription_id = azure_subscription_id,
        resource_group = resource_group,
        source_capacity = '', # The name of the Trial capacity.
        target_capacity = '', # The name of Fabric capacity (if it does not exist it will be created).
        target_capacity_sku = "F64", # Defaults to F64 but can specify the required SKU.
        target_capacity_admin_members = None, # Setting this to None will use the same admin members as the Trial capacity.
    )