#! /usr/bin/env python3
from azure.mgmt.storage import StorageManagementClient
from azure.mgmt.storage.models import StorageAccountUpdateParameters, Sku, SkuName


def list_storage_accounts(credentials, rgs):
    print("Fetching storage accounts data")
    storage_account_data = {}
    for sub, groups in rgs.items():
        storage_client = StorageManagementClient(credentials, sub)
        for rg in groups:
            storage_accounts = storage_client.storage_accounts.list_by_resource_group(
                rg
            )
            for account in storage_accounts:
                account_data = dict(
                    {
                        "name": account.name,
                        "sub": sub,
                        "resource_group": rg,
                        "sku": account.sku.name,
                    }
                )
                storage_account_data[account.name] = account_data

    return storage_account_data


def change_sku(credentials, storage_accounts):
    storage_container_data = {}
    for account, config in storage_accounts.items():
        storage_client = StorageManagementClient(credentials, config.get("sub"))
        storage_container_data[account] = account
        if (str(config.get("sku"))) == "SkuName.standard_lrs":
            storage_client.storage_accounts.update(
                config.get("resource_group"),
                account,
                StorageAccountUpdateParameters(sku=Sku(name=SkuName.standard_grs)),
            )
            print(f"Storage account {account} changed to standard_grs SKU")


def change_storage_sku(credentials, rgs):
    storage_accounts = list_storage_accounts(credentials, rgs)
    change_sku(credentials, storage_accounts)