from pathlib import Path

import pytest
import yaml
from marshmallow.exceptions import ValidationError

from azure.ai.ml._schema.registry import RegistrySchema
from azure.ai.ml.constants._common import BASE_PATH_CONTEXT_KEY, PublicNetworkAccess
from azure.ai.ml.constants._registry import AcrAccountSku, StorageAccountType
from azure.ai.ml.entities import RegistryRegionArmDetails, SystemCreatedAcrAccount, SystemCreatedStorageAccount
from azure.ai.ml.entities._util import load_from_dict


@pytest.mark.unittest
class TestRegistrySchema:
    def test_deserialize_from_yaml(self) -> None:
        path = Path("./tests/test_configs/registry/registry_valid.yaml")
        with open(path, "r") as f:
            target = yaml.safe_load(f)
            context = {BASE_PATH_CONTEXT_KEY: path.parent}
            registry = load_from_dict(RegistrySchema, target, context)
            assert registry
            assert registry["description"] == "This is a registry description"
            assert registry["name"] == "registry_name"
            assert registry["id"] == "registry_id"
            assert registry["tags"] == {"purpose": "testing", "other_tag": "value"}
            assert registry["location"] == "EastUS2"
            assert registry["public_network_access"] == PublicNetworkAccess.DISABLED
            assert registry["intellectual_property_publisher"] == "registry_publisher"
            assert (
                registry["container_registry"]
                == "/subscriptions/sub_id/resourceGroups/some_rg/providers/Microsoft.ContainerRegistry/registries/acr_id"
            )
            assert len(registry["replication_locations"]) == 1

            detail = registry["replication_locations"][0]
            assert isinstance(detail, RegistryRegionArmDetails)
            assert detail.location == "EastUS"
            storages = detail.storage_config
            assert len(storages) == 2

            assert (
                storages[0]
                == "/subscriptions/sub_id/resourceGroups/some_rg/providers/Microsoft.Storage/storageAccounts/some_storage_account"
            )
            assert isinstance(storages[1], SystemCreatedStorageAccount)
            assert not storages[1].storage_account_hns
            assert storages[1].storage_account_type == StorageAccountType.STANDARD_RAGRS

    def test_deserialize_from_yaml_with_system_acr(self) -> None:
        path = Path("./tests/test_configs/registry/registry_valid_2.yaml")
        with open(path, "r") as f:
            target = yaml.safe_load(f)
            context = {BASE_PATH_CONTEXT_KEY: path.parent}
            registry = load_from_dict(RegistrySchema, target, context)
            assert registry
            assert isinstance(registry["container_registry"], SystemCreatedAcrAccount)
            assert registry["container_registry"].acr_account_sku == AcrAccountSku.PREMIUM

    def test_deserialize_from_yaml_with_no_acr(self) -> None:
        path = Path("./tests/test_configs/registry/registry_valid_3.yaml")
        with open(path, "r") as f:
            target = yaml.safe_load(f)
            context = {BASE_PATH_CONTEXT_KEY: path.parent}
            registry = load_from_dict(RegistrySchema, target, context)
            assert registry
            assert isinstance(registry["container_registry"], SystemCreatedAcrAccount)
            assert registry["container_registry"].acr_account_sku == AcrAccountSku.PREMIUM

    def test_deserialize_bad_storage_account_type(self) -> None:
        path = Path("./tests/test_configs/registry/registry_bad_storage_account_type.yaml")
        with open(path, "r") as f:
            target = yaml.safe_load(f)
            context = {BASE_PATH_CONTEXT_KEY: path.parent}
            with pytest.raises(Exception) as e_info:
                load_from_dict(RegistrySchema, target, context)
            assert e_info
            assert isinstance(e_info._excinfo[1], ValidationError)
            assert "Value 'NOT_A_REAL_ACCOUNT_TYPE' passed is not in set" in e_info._excinfo[1].messages[0]

    def test_deserialize_bad_arm_resource_id(self) -> None:
        path = Path("./tests/test_configs/registry/registry_bad_arm_resource_id.yaml")
        with open(path, "r") as f:
            target = yaml.safe_load(f)
            context = {BASE_PATH_CONTEXT_KEY: path.parent}
            with pytest.raises(Exception) as e_info:
                load_from_dict(RegistrySchema, target, context)
            assert e_info
            assert isinstance(e_info._excinfo[1], ValidationError)
            assert "container_registry" in e_info._excinfo[1].messages[0]
            assert "Invalid value" in e_info._excinfo[1].messages[0]
            # import pdb; pdb.set_trace()
