"""Upgrade command for handling breaking changes when updating ACA-PY versions."""

import asyncio
import yaml

from configargparse import ArgumentParser
from packaging import version as package_version
from typing import Callable, Sequence, Optional

from ..core.profile import Profile
from ..config import argparse as arg
from ..config.default_context import DefaultContextBuilder
from ..config.base import BaseError
from ..config.util import common_config
from ..config.wallet import wallet_config
from ..messaging.models.base_record import BaseRecord
from ..storage.base import BaseStorage
from ..storage.error import StorageNotFoundError
from ..storage.record import StorageRecord
from ..utils.classloader import ClassLoader, ClassNotFoundError
from ..version import __version__, RECORD_TYPE_ACAPY_VERSION

from . import PROG

DEFAULT_UPGRADE_CONFIG_PATH = (
    "./aries_cloudagent/commands/default_version_upgrade_config.yml"
)


class UpgradeError(BaseError):
    """Base exception for upgrade related errors."""


class VersionUpgradeConfig:
    """Handle ACA-Py version upgrade config."""

    def __init__(self, config_path: str = None):
        """Initialize config for use during upgrade process."""
        self.function_map_config = {}
        self.upgrade_configs = {}
        self.setup_executable_map_config(CONFIG_v7_3)
        if config_path:
            self.setup_version_upgrade_config(config_path)
        else:
            self.setup_version_upgrade_config(DEFAULT_UPGRADE_CONFIG_PATH)

    def setup_executable_map_config(self, config_dict: dict):
        """Set ups config with reference to functions mapped to versions."""
        for version, config in config_dict.items():
            self.function_map_config[version] = {}
            if "update_existing_function_inst" in config:
                self.function_map_config[version][
                    "update_existing_function_inst"
                ] = config.get("update_existing_function_inst")

    def setup_version_upgrade_config(self, path: str):
        """Set ups config dict from the provided YML file."""
        with open(path, "r") as stream:
            config_dict = yaml.safe_load(stream)
            version_config_dict = {}
            for version, provided_config in config_dict.items():
                recs_list = []
                version_config_dict[version] = {}
                if "resave_records" in provided_config:
                    if provided_config.get("resave_records").get("base_record_path"):
                        recs_list = recs_list + provided_config.get(
                            "resave_records"
                        ).get("base_record_path")
                    if provided_config.get("resave_records").get(
                        "base_exch_record_path"
                    ):
                        recs_list = recs_list + provided_config.get(
                            "resave_records"
                        ).get("base_exch_record_path")
                version_config_dict[version]["resave_records"] = recs_list
                version_config_dict[version]["update_existing_records"] = (
                    provided_config.get("update_existing_records") or False
                )
            if version_config_dict == {}:
                raise UpgradeError(f"No version configs found in {path}")
            self.upgrade_configs = version_config_dict

    def get_update_existing_func(self, ver: str) -> Optional[Callable]:
        """Return callable update_existing_records function for specific version."""
        if ver in self.function_map_config:
            return self.function_map_config.get(ver).get(
                "update_existing_function_inst"
            )
        else:
            return None


def init_argument_parser(parser: ArgumentParser):
    """Initialize an argument parser with the module's arguments."""
    return arg.load_argument_groups(parser, *arg.group.get_registered(arg.CAT_UPGRADE))


async def upgrade(settings: dict):
    """Perform upgradation steps."""
    context_builder = DefaultContextBuilder(settings)
    context = await context_builder.build_context()
    try:
        version_upgrade_config_inst = VersionUpgradeConfig(
            settings.get("upgrade.config_path")
        )
        upgrade_configs = version_upgrade_config_inst.upgrade_configs
        root_profile, public_did = await wallet_config(context)
        version_storage_record = None
        upgrade_to_version = f"v{__version__}"
        versions_found_in_config = upgrade_configs.keys()
        sorted_versions_found_in_config = sorted(
            versions_found_in_config, key=lambda x: package_version.parse(x)
        )
        async with root_profile.session() as session:
            storage = session.inject(BaseStorage)
            try:
                version_storage_record = await storage.find_record(
                    type_filter=RECORD_TYPE_ACAPY_VERSION, tag_query={}
                )
                upgrade_from_version = version_storage_record.value
                if "upgrade.from_version" in settings:
                    print(
                        (
                            f"version {upgrade_from_version} found in storage"
                            ", --from-version will be ignored."
                        )
                    )
            except StorageNotFoundError:
                if "upgrade.from_version" in settings:
                    upgrade_from_version = settings.get("upgrade.from_version")
                else:
                    upgrade_from_version = sorted_versions_found_in_config[-1]
                    print(
                        "No ACA-Py version found in wallet storage and "
                        "no --from-version specified. Selecting "
                        f"{upgrade_from_version} as --from-version from "
                        "the config."
                    )
        if upgrade_from_version == upgrade_to_version:
            raise UpgradeError(
                f"Version {upgrade_from_version} to upgrade from and "
                f"current version to upgrade to {upgrade_to_version} "
                "are same."
            )
        if upgrade_from_version not in sorted_versions_found_in_config:
            raise UpgradeError(
                f"No upgrade configuration found for {upgrade_from_version}"
            )
        upgrade_from_version_index = sorted_versions_found_in_config.index(
            upgrade_from_version
        )
        for config_from_version in sorted_versions_found_in_config[
            upgrade_from_version_index:
        ]:
            print(f"Running upgrade process for {config_from_version}")
            upgrade_config = upgrade_configs.get(config_from_version)
            # Step 1 re-saving all BaseRecord and BaseExchangeRecord
            if "resave_records" in upgrade_config:
                resave_record_paths = upgrade_config.get("resave_records")
                for record_path in resave_record_paths:
                    try:
                        record_type = ClassLoader.load_class(record_path)
                    except ClassNotFoundError as err:
                        raise UpgradeError(
                            f"Unknown Record type {record_path}"
                        ) from err
                    if not issubclass(record_type, BaseRecord):
                        raise UpgradeError(
                            f"Only BaseRecord can be resaved, found: {str(record_type)}"
                        )
                    async with root_profile.session() as session:
                        all_records = await record_type.query(session)
                        for record in all_records:
                            await record.save(
                                session,
                                reason="re-saving record during ACA-Py upgrade process",
                            )
                        if len(all_records) == 0:
                            print(f"No records of {str(record_type)} found")
                        else:
                            print(
                                f"All records of {str(record_type)} successfully re-saved"
                            )
            # Step 2 Update existing records, if required
            if (
                "update_existing_records" in upgrade_config
                and upgrade_config.get("update_existing_records") is True
            ):
                update_existing_recs_callable = (
                    version_upgrade_config_inst.get_update_existing_func(
                        config_from_version
                    )
                )
                if not update_existing_recs_callable:
                    raise UpgradeError(
                        "No update_existing_records function "
                        f"specified for {config_from_version}"
                    )
                await update_existing_recs_callable(root_profile)
        # Update storage version
        async with root_profile.session() as session:
            storage = session.inject(BaseStorage)
            if not version_storage_record:
                await storage.add_record(
                    StorageRecord(
                        RECORD_TYPE_ACAPY_VERSION,
                        upgrade_to_version,
                    )
                )
            else:
                await storage.update_record(
                    version_storage_record, upgrade_to_version, {}
                )
        await root_profile.close()
    except BaseError as e:
        raise UpgradeError(f"Error during upgrade: {e}")


async def update_existing_records(profile: Profile):
    """
    Update existing records.

    Args:
        profile: Root profile

    """
    pass


def execute(argv: Sequence[str] = None):
    """Entrypoint."""
    parser = arg.create_argument_parser(prog=PROG)
    parser.prog += " upgrade"
    get_settings = init_argument_parser(parser)
    args = parser.parse_args(argv)
    settings = get_settings(args)
    common_config(settings)
    loop = asyncio.get_event_loop()
    loop.run_until_complete(upgrade(settings))


def main():
    """Execute the main line."""
    if __name__ == "__main__":
        execute()


# Update every release
CONFIG_v7_3 = {
    "v0.7.2": {
        "update_existing_function_inst": update_existing_records,
    },
}


main()

"""Upgrade command for handling breaking changes when updating ACA-PY versions."""

import asyncio
import yaml

from configargparse import ArgumentParser
from packaging import version as package_version
from typing import Callable, Sequence, Optional

from ..core.profile import Profile
from ..config import argparse as arg
from ..config.default_context import DefaultContextBuilder
from ..config.base import BaseError
from ..config.util import common_config
from ..config.wallet import wallet_config
from ..messaging.models.base_record import BaseRecord
from ..storage.base import BaseStorage
from ..storage.error import StorageNotFoundError
from ..storage.record import StorageRecord
from ..utils.classloader import ClassLoader, ClassNotFoundError
from ..version import __version__, RECORD_TYPE_ACAPY_VERSION

from . import PROG

DEFAULT_UPGRADE_CONFIG_PATH = (
    "./aries_cloudagent/commands/default_version_upgrade_config.yml"
)


class UpgradeError(BaseError):
    """Base exception for upgrade related errors."""


class VersionUpgradeConfig:
    """Handle ACA-Py version upgrade config."""

    def __init__(self, config_path: str = None):
        """Initialize config for use during upgrade process."""
        self.function_map_config = {}
        self.upgrade_configs = {}
        self.setup_executable_map_config(CONFIG_v7_3)
        if config_path:
            self.setup_version_upgrade_config(config_path)
        else:
            self.setup_version_upgrade_config(DEFAULT_UPGRADE_CONFIG_PATH)

    def setup_executable_map_config(self, config_dict: dict):
        """Set ups config with reference to functions mapped to versions."""
        for version, config in config_dict.items():
            self.function_map_config[version] = {}
            if "update_existing_function_inst" in config:
                self.function_map_config[version][
                    "update_existing_function_inst"
                ] = config.get("update_existing_function_inst")

    def setup_version_upgrade_config(self, path: str):
        """Set ups config dict from the provided YML file."""
        with open(path, "r") as stream:
            config_dict = yaml.safe_load(stream)
            version_config_dict = {}
            for version, provided_config in config_dict.items():
                recs_list = []
                version_config_dict[version] = {}
                if "resave_records" in provided_config:
                    if provided_config.get("resave_records").get("base_record_path"):
                        recs_list = recs_list + provided_config.get(
                            "resave_records"
                        ).get("base_record_path")
                    if provided_config.get("resave_records").get(
                        "base_exch_record_path"
                    ):
                        recs_list = recs_list + provided_config.get(
                            "resave_records"
                        ).get("base_exch_record_path")
                version_config_dict[version]["resave_records"] = recs_list
                version_config_dict[version]["update_existing_records"] = (
                    provided_config.get("update_existing_records") or False
                )
            if version_config_dict == {}:
                raise UpgradeError(f"No version configs found in {path}")
            self.upgrade_configs = version_config_dict

    def get_update_existing_func(self, ver: str) -> Optional[Callable]:
        """Return callable update_existing_records function for specific version."""
        if ver in self.function_map_config:
            return self.function_map_config.get(ver).get(
                "update_existing_function_inst"
            )
        else:
            return None


def init_argument_parser(parser: ArgumentParser):
    """Initialize an argument parser with the module's arguments."""
    return arg.load_argument_groups(parser, *arg.group.get_registered(arg.CAT_UPGRADE))


async def upgrade(settings: dict):
    """Perform upgradation steps."""
    context_builder = DefaultContextBuilder(settings)
    context = await context_builder.build_context()
    try:
        version_upgrade_config_inst = VersionUpgradeConfig(
            settings.get("upgrade.config_path")
        )
        upgrade_configs = version_upgrade_config_inst.upgrade_configs
        root_profile, public_did = await wallet_config(context)
        version_storage_record = None
        upgrade_to_version = f"v{__version__}"
        versions_found_in_config = upgrade_configs.keys()
        sorted_versions_found_in_config = sorted(
            versions_found_in_config, key=lambda x: package_version.parse(x)
        )
        async with root_profile.session() as session:
            storage = session.inject(BaseStorage)
            try:
                version_storage_record = await storage.find_record(
                    type_filter=RECORD_TYPE_ACAPY_VERSION, tag_query={}
                )
                upgrade_from_version = version_storage_record.value
                if "upgrade.from_version" in settings:
                    print(
                        (
                            f"version {upgrade_from_version} found in storage"
                            ", --from-version will be ignored."
                        )
                    )
            except StorageNotFoundError:
                if "upgrade.from_version" in settings:
                    upgrade_from_version = settings.get("upgrade.from_version")
                else:
                    upgrade_from_version = sorted_versions_found_in_config[-1]
                    print(
                        "No ACA-Py version found in wallet storage and "
                        "no --from-version specified. Selecting "
                        f"{upgrade_from_version} as --from-version from "
                        "the config."
                    )
        if upgrade_from_version == upgrade_to_version:
            raise UpgradeError(
                f"Version {upgrade_from_version} to upgrade from and "
                f"current version to upgrade to {upgrade_to_version} "
                "are same."
            )
        if upgrade_from_version not in sorted_versions_found_in_config:
            raise UpgradeError(
                f"No upgrade configuration found for {upgrade_from_version}"
            )
        upgrade_from_version_index = sorted_versions_found_in_config.index(
            upgrade_from_version
        )
        for config_from_version in sorted_versions_found_in_config[
            upgrade_from_version_index:
        ]:
            print(f"Running upgrade process for {config_from_version}")
            upgrade_config = upgrade_configs.get(config_from_version)
            # Step 1 re-saving all BaseRecord and BaseExchangeRecord
            if "resave_records" in upgrade_config:
                resave_record_paths = upgrade_config.get("resave_records")
                for record_path in resave_record_paths:
                    try:
                        record_type = ClassLoader.load_class(record_path)
                    except ClassNotFoundError as err:
                        raise UpgradeError(
                            f"Unknown Record type {record_path}"
                        ) from err
                    if not issubclass(record_type, BaseRecord):
                        raise UpgradeError(
                            f"Only BaseRecord can be resaved, found: {str(record_type)}"
                        )
                    async with root_profile.session() as session:
                        all_records = await record_type.query(session)
                        for record in all_records:
                            await record.save(
                                session,
                                reason="re-saving record during ACA-Py upgrade process",
                            )
                        if len(all_records) == 0:
                            print(f"No records of {str(record_type)} found")
                        else:
                            print(
                                f"All records of {str(record_type)} successfully re-saved"
                            )
            # Step 2 Update existing records, if required
            if (
                "update_existing_records" in upgrade_config
                and upgrade_config.get("update_existing_records") is True
            ):
                update_existing_recs_callable = (
                    version_upgrade_config_inst.get_update_existing_func(
                        config_from_version
                    )
                )
                if not update_existing_recs_callable:
                    raise UpgradeError(
                        "No update_existing_records function "
                        f"specified for {config_from_version}"
                    )
                await update_existing_recs_callable(root_profile)
        # Update storage version
        async with root_profile.session() as session:
            storage = session.inject(BaseStorage)
            if not version_storage_record:
                await storage.add_record(
                    StorageRecord(
                        RECORD_TYPE_ACAPY_VERSION,
                        upgrade_to_version,
                    )
                )
            else:
                await storage.update_record(
                    version_storage_record, upgrade_to_version, {}
                )
        await root_profile.close()
    except BaseError as e:
        raise UpgradeError(f"Error during upgrade: {e}")


async def update_existing_records(profile: Profile):
    """
    Update existing records.

    Args:
        profile: Root profile

    """
    pass


def execute(argv: Sequence[str] = None):
    """Entrypoint."""
    parser = arg.create_argument_parser(prog=PROG)
    parser.prog += " upgrade"
    get_settings = init_argument_parser(parser)
    args = parser.parse_args(argv)
    settings = get_settings(args)
    common_config(settings)
    loop = asyncio.get_event_loop()
    loop.run_until_complete(upgrade(settings))


def main():
    """Execute the main line."""
    if __name__ == "__main__":
        execute()


# Update every release
CONFIG_v7_3 = {
    "v0.7.2": {
        "update_existing_function_inst": update_existing_records,
    },
}


main()

Upgrade command for handling breaking changes when updating ACA-PY versions. Base exception for upgrade related errors. Handle ACA-Py version upgrade config. Initialize config for use during upgrade process. Set ups config with reference to functions mapped to versions. Set ups config dict from the provided YML file. Return callable update_existing_records function for specific version. Initialize an argument parser with the module's arguments. Perform upgradation steps. # Step 1 re-saving all BaseRecord and BaseExchangeRecord # Step 2 Update existing records, if required # Update storage version Update existing records.      Args:         profile: Root profile Entrypoint. Execute the main line. # Update every release