#
# Copyright (c) 2020 by Delphix. All rights reserved.
#

from __future__ import print_function
from __future__ import unicode_literals

import argparse
import collections
import sys

import delphixpy.v1_11_3
from delphixpy.v1_11_3.web import vo
from delphixpy.v1_11_3.web.source import operationTemplate

SCRIPT_DESCRIPTION = """
%(prog)s:
    This script is used to migrate MSSQL dSource and VDB hooks and
    hook templates to the target host's default PowerShell version.
    Afterwards the migration, the script disables the "INSTALLEDPOWERSHELL"
    feature flag, if it is enabled.

    By default, both hooks and hook templates will be migrated. It
    is possible to migrate only the hooks or only the hook templates
    by setting the  "--migrate-only" optional param to "hooks" or
    "templates" respectively.

    By setting the value of the optional param- "--hook-ps-version"
    to "ps2", hooks can be migrated to PowerShell version 2 as well.
    Similarly, by setting the value of optional param-
    '--hook-templ-ps-version' to 'ps2', hook templates can also be
    migrated to PowerShell version 2.
    """


FEATURE_FLAG_INSTALLEDPOWERSHELL = "INSTALLEDPOWERSHELL"

# Every valid hook operation for an MSSql linked source.
VALID_LINKED_SRC_OPS = ["pre_sync", "post_sync"]

# Every valid hook operation for an MSSql virtual source.
VALID_VIRTUAL_SRC_OPS = [
    "configure_clone",
    "pre_refresh",
    "post_refresh",
    "pre_rollback",
    "post_rollback",
    "pre_snapshot",
    "post_snapshot",
    "pre_start",
    "post_start",
    "pre_stop",
    "post_stop",
]


MSsqlHookType = collections.namedtuple(
    "MSsqlHookType", ["name", "powershell_version"]
)


class MSSqlHookTypes(object):
    """
    Container for storing the name of the type of hook and hook template,
    and the PowerShell version it corresponds to.
    """

    RunDefaultPowerShellOnSourceOperation = MSsqlHookType(
        "RunDefaultPowerShellOnSourceOperation", "Default PowerShell Version"
    )
    RunPowerShellOnSourceOperation = MSsqlHookType(
        "RunPowerShellOnSourceOperation", "PowerShell Version 2"
    )

    @classmethod
    def all(cls):
        return [
            cls.RunDefaultPowerShellOnSourceOperation,
            cls.RunPowerShellOnSourceOperation,
        ]


class InputsPowerShellVersionArgument(object):
    """
    Container for storing the options the user of the script can specify
    for indicating the PowerShell version to which the hooks should
    be migrated.
    """

    PS2 = "ps2"
    DEFAULT = "default"

    @classmethod
    def all(cls):
        return [cls.DEFAULT, cls.PS2]


class InputsMigrateOnlyArgument(object):
    """
    Container for storing the options the user of the script can specify
    for `--migrate-only` argument, via the `--migrate-only` argument the
    user can choose to migrate only the hooks or the hook templates.
    """

    HOOKS = "hooks"
    HOOK_TEMPLATES = "templates"
    HOOKS_AND_TEMPLATES = "hooks-and-templates"

    @classmethod
    def all(cls):
        return [cls.HOOKS, cls.HOOK_TEMPLATES, cls.HOOKS_AND_TEMPLATES]


class DelphixEngineConnectivityError(Exception):
    """
    Raised when the host on which the script is being run doesn't have
    connectivity with the provided Delphix engine on which the hooks migration
    needs to be performed, with the given command line arguments.
    """


def main():
    """
    The function receives the command line arguments via the function-
    `create_argument_parser` and uses the `1.11.3` APIs from the
    `delphixpy.v1_11_3` package to create a Delphix engine object and then
    executes the functions- `migrate_hook` and `migrate_hook_template`;
    if applicable `disable_feature_flag` function is also run.
    """
    parser = create_argument_parser()
    args = parser.parse_args()
    dlpx_host_address = args.engine_addr
    dlpx_sa_user = args.sys_admin_usr
    dlpx_sa_password = args.sys_admin_pwd
    dlpx_admin_user = args.admin_usr
    dlpx_admin_password = args.admin_pwd
    hook_ps_version = args.hook_ps_version
    hook_templ_ps_version = args.hook_templ_ps_version
    migrate_only = args.migrate_only

    try:
        engine_sa = delphixpy.v1_11_3.delphix_engine.DelphixEngine(
            dlpx_host_address, dlpx_sa_user, dlpx_sa_password, "SYSTEM"
        )
        engine_admin = delphixpy.v1_11_3.delphix_engine.DelphixEngine(
            dlpx_host_address, dlpx_admin_user, dlpx_admin_password, "DOMAIN"
        )

        # Sanity check that the engine is accessible from the given host with
        # the arguments/inputs provided.
        check_engine_connectivity(dlpx_host_address, engine_sa, engine_admin)

        # Run the migration.
        hook_type = _get_hook_type(hook_ps_version)
        hook_templ_type = _get_hook_type(hook_templ_ps_version)
        if migrate_only == InputsMigrateOnlyArgument.HOOKS:
            migrate_hook(engine_admin, hook_type)
        elif migrate_only == InputsMigrateOnlyArgument.HOOK_TEMPLATES:
            migrate_hook_template(engine_admin, hook_templ_type)
        else:
            migrate_hook(engine_admin, hook_type)
            migrate_hook_template(engine_admin, hook_templ_type)

        # Disable the feature flag if it is enabled.
        enabled_feature_flags = delphixpy.v1_11_3.web.about.get(
            engine_sa
        ).enabled_features
        if FEATURE_FLAG_INSTALLEDPOWERSHELL in enabled_feature_flags:
            disable_feature_flag(engine_sa, dlpx_host_address)

    except Exception:
        if not args.debug:
            sys.tracebacklimit = 0
        raise

    sys.exit(0)


def create_argument_parser():
    """
    The function receives the command line arguments for the script.

    :returns: The command line arguments received from the user for the script
        params.
    :rtype: :py:class: `argparse.ArgumentParser`
    """
    parser = argparse.ArgumentParser(
        description=SCRIPT_DESCRIPTION,
        formatter_class=argparse.RawTextHelpFormatter,
        add_help=False,
    )
    required = parser.add_argument_group("Required arguments")
    optional = parser.add_argument_group("Optional arguments")

    required.add_argument(
        "--engine-addr", help="The Delphix Engine host address.", required=True
    )
    required.add_argument(
        "--sys-admin-usr",
        help=(
            "The username for System Administrator user to log into the"
            " Delphix Engine.\n"
            "System Administrator login is required to disable the feature"
            " flag."
        ),
        required=True,
    )
    required.add_argument(
        "--sys-admin-pwd",
        help=(
            "The password for the System Administrator user to log into the"
            " Delphix Engine.\n"
            "System Administrator login is required to disable the feature"
            " flag."
        ),
        required=True,
    )
    required.add_argument(
        "--admin-usr",
        help=(
            "The username for Delphix Administrator user to log into the"
            " Delphix Engine.\n"
            "Delphix Administrator login is required to update the hooks."
        ),
        required=True,
    )
    required.add_argument(
        "--admin-pwd",
        help=(
            "The password for the Delphix Administrator user to log into the"
            " Delphix Engine.\n"
            "Delphix Administrator login is required to update the hooks."
        ),
        required=True,
    )

    optional.add_argument(
        "-h",
        "--help",
        action="help",
        help="Show this help message and exit.\n\n",
    )
    optional.add_argument(
        "--debug",
        action="store_true",
        help=(
            "Run the script in debug mode.\n"
            "With this argument, in case of an error, the Python"
            " stack trace for the error will be printed.\n\n"
        ),
    )
    help_menu_ps_version_arg = (
        "The PowerShell version to which the {} should be updated:\n"
        "A. Input 'default' for default PowerShell version.\n"
        "B. Input 'ps2' for PowerShell version 2.\n"
        "Default value is '%(default)s'.\n\n"
    )
    optional.add_argument(
        "--hook-ps-version",
        help=help_menu_ps_version_arg.format("hooks"),
        choices=InputsPowerShellVersionArgument.all(),
        default=InputsPowerShellVersionArgument.DEFAULT,
    )
    optional.add_argument(
        "--hook-templ-ps-version",
        help=help_menu_ps_version_arg.format("hook templates"),
        choices=InputsPowerShellVersionArgument.all(),
        default=InputsPowerShellVersionArgument.DEFAULT,
    )
    optional.add_argument(
        "--migrate-only",
        help="Choose whether to migrate only the hooks or only the hook"
        " templates.\n"
        "A. Input 'hooks' to migrate only the hooks.\n"
        "B. Input 'templates' to migrate only the hook templates.\n"
        "Default value is '%(default)s', which will migrate both"
        " hooks and hook templates.\n\n",
        choices=InputsMigrateOnlyArgument.all(),
        default="hooks-and-templates",
    )

    return parser


def check_engine_connectivity(dlpx_host_address, engine_sa, engine_admin):
    """
    The function performs a sanity check that the engine is accessible
    from the host on which the script has been run, and with the
    arguments/inputs provided by making an attempt to login to the Delphix
    engine via given System Administrator and Delphix Administrator users.

    :param dlpx_host_address: The Delphix engine host address.
    :type dlpx_host_address: ``str``
    :param engine_sa: Delphix engine object using System Administrator user.
    :type engine_sa:
        :py:class: `delphixpy.v1_11_3.delphix_engine.DelphixEngine`
    :param engine_admin: Delphix engine object using Delphix Administrator
        user.
    :type engine_admin:
        :py:class: `delphixpy.v1_11_3.delphix_engine.DelphixEngine`
    :raises: :py:class: `DelphixEngineConnectivityError` if the given Delphix
        engine can not be accessed via the host, or with the given command
        line arguments received.
    """
    print(
        "\n\nPerforming connectivity test between this host on which"
        " the script has been run and the given Delphix engine-"
        " [{engine}] using the given System and Admin users ...".format(
            engine=dlpx_host_address
        )
    )

    engine_sa_login_err = None
    try:
        engine_sa.login()
    except Exception as e:
        engine_sa_login_err = str(e)
        print(
            "\nError occurred while connecting to the engine via the given"
            " Delphix System Administrator user: \n{}\n".format(
                engine_sa_login_err
            )
        )

    engine_admin_login_err = None
    try:
        engine_admin.login()
    except Exception as e:
        engine_admin_login_err = str(e)
        print(
            "\nError occurred while connecting to the engine via the given"
            " Delphix Administrator user: \n{}\n".format(
                engine_admin_login_err
            )
        )

    if any([engine_sa_login_err, engine_admin_login_err]):
        err_message = (
            "\nThe connectivity test between the host and the Delphix"
            " engine failed.\nPlease check if the Delphix engine is"
            " accessible from this host, and the Delphix System and Admin"
            " user input arguments provided to the script are correct.\n"
        )
        raise DelphixEngineConnectivityError(err_message)
    else:
        print(
            "\nThe connectivity test between the host and the Delphix"
            " engine passed."
        )


def _get_hook_type(input_ps_version):
    """
    Returns the type to which the hooks or hook templates should
    be migrated based on the input argument value.

    :param input_ps_version: The option the user of the script specified
        for indicating the PowerShell version to which the hooks or the hook
        templates should be migrated.
    :type input_ps_version: ``str``
    :returns: Returns the type to which the hooks or hook templates should
        be migrated.
    :rtype: :py:class: :py:func: `collections.namedtuple`
    """
    if input_ps_version == InputsPowerShellVersionArgument.DEFAULT:
        return MSSqlHookTypes.RunDefaultPowerShellOnSourceOperation
    elif input_ps_version == InputsPowerShellVersionArgument.PS2:
        return MSSqlHookTypes.RunPowerShellOnSourceOperation


def migrate_hook(engine_obj, hook_type):
    """
    The function migrates the hooks for the MSSql dSources and VDBs in the
    given engine object to the type specified in `hook_type` param.

    :param engine_obj: The Delphix engine object.
    :type engine_obj:
        :py:class:`delphixpy.v1_11_3.delphix_engine.DelphixEngine`
    :param hook_type: The MSSql hook type to which the hooks needs to
        be migrated.
    :type hook_type: :py:func: `collections.namedtuple`
    """

    print(
        "\nUpdating the PowerShell version for the MSSql Database"
        " hooks to- [{ps_ver}] ...\n".format(
            ps_ver=hook_type.powershell_version
        )
    )
    all_databases = delphixpy.v1_11_3.web.database.get_all(engine_obj)
    for db in all_databases:
        if isinstance(db, vo.MSSqlDatabaseContainer):
            migrate_hook_for_database(engine_obj, db, hook_type.name)
    print(
        "\nSuccessfully updated the PowerShell version for the MSSql Database"
        " hooks to- [{ps_ver}].\n".format(ps_ver=hook_type.powershell_version)
    )


def migrate_hook_template(engine_admin, hook_templ_type):
    """
    The function migrates the MSSql hook templates in the given engine object
    to the type specified in `hook_type` param.

    :param engine_admin: The Delphix engine object.
    :type engine_admin:
        :py:class:`delphixpy.v1_11_3.delphix_engine.DelphixEngine`
    :param hook_templ_type: The MSSql hook type to which the hook templates
        needs to be migrated.
    :type hook_templ_type: :py:func: `collections.namedtuple`
    """

    print(
        "\nUpdating the PowerShell version for the MSSql Database hook"
        " templates to- [{ps_ver}] ...\n".format(
            ps_ver=hook_templ_type.powershell_version
        )
    )
    hook_type_names = [
        hook_templ_type_.name for hook_templ_type_ in MSSqlHookTypes.all()
    ]
    all_hook_templs = operationTemplate.get_all(engine_admin)
    for hook_template in all_hook_templs:
        # Skip if it is not an MSSQL hook template.
        if hook_template.operation.type not in hook_type_names:
            continue

        # Skip if the hook template is already of the desired type.
        if hook_template.operation.type == hook_templ_type.name:
            continue

        # Perform the update.
        operation = getattr(hook_template, "operation")
        hook_template.operation.type = hook_templ_type.name
        hook_template.operation = operation
        operationTemplate.update(
            engine_admin, hook_template.reference, hook_template
        )

        print(
            "PowerShell version of the Hook template- [{hook_templ_name}]"
            " has been updated.".format(hook_templ_name=hook_template.name)
        )
    print(
        "\nSuccessfully updated the PowerShell version for the MSSql Database"
        " hook templates to- [{ps_ver}].\n".format(
            ps_ver=hook_templ_type.powershell_version
        )
    )


def migrate_hook_for_database(engine_obj, database_, hook_type):
    """
    The function migrates all the hooks for the given MSSql dSource or VDB
    to the type specified in `hook_type` param.

    :param engine_obj: The Delphix engine object.
    :type engine_obj:
        :py:class:`delphixpy.v1_11_3.delphix_engine.DelphixEngine`
    :param database_: The MSSql dSource or VDB for which the hooks needs to
        be migrated.
    :type database_:
        :py:class:`delphixpy.v1_11_3.web.objects.MSSqlDatabaseContainer`
    :param hook_type: The name of the MSSql hook type to which the hooks needs
        to be migrated.
    :type hook_type: ``str``
    """
    hooks_present = []
    database_src = delphixpy.v1_11_3.web.source.get_all(
        engine_obj, database=database_.reference
    )[0]
    if isinstance(database_src, vo.MSSqlLinkedSource):
        link_src_ops_obj = vo.LinkedSourceOperations()
        for valid_op in VALID_LINKED_SRC_OPS:
            hooks_present.append(
                _set_attr_for_src_ops_obj(
                    database_src, link_src_ops_obj, valid_op, hook_type
                )
            )
        database_src.operations = link_src_ops_obj
    elif isinstance(database_src, vo.MSSqlVirtualSource):
        virtual_src_ops_obj = vo.VirtualSourceOperations()
        for valid_op in VALID_VIRTUAL_SRC_OPS:
            hooks_present.append(
                _set_attr_for_src_ops_obj(
                    database_src, virtual_src_ops_obj, valid_op, hook_type
                )
            )
        database_src.operations = virtual_src_ops_obj

    # Update the database only if hook(s) are present.
    if any(hooks_present):
        delphixpy.v1_11_3.web.source.update(
            engine_obj, database_src.reference, database_src
        )
        print(
            "PowerShell version of the Hooks for the database-"
            " [{db}] has been updated.".format(db=database_.name)
        )


def _set_attr_for_src_ops_obj(
    database_src, src_operations, valid_op, hook_type
):
    """
    The function sets the attributes for the given `LinkedSourceOperations`
    or `VirtualSourceOperations` object, this object will be set as the
    `operations` attribute for the given database source object.

    For the same, the function first checks whether the given database source
    object(MSSqlLinkedSource or MSSqlVirtualSource) has the given operation
    (`valid_op`) as an attribute in its `operations` attribute, if so, it
    retrieves all the hooks for the operation and updates the `type` attribute
    for each hook in the operation type- `valid_op` to the value specified
    in `hook_type`.

    :param database_src: The source object for MSSql dSource or VDB for which
        the hooks needs to be migrated.
    :type database_src:
        :py:class:`delphixpy.v1_11_3.web.objects.MSSqlLinkedSource`
        or :py:class:`delphixpy.v1_11_3.web.objects.MSSqlVirtualSource`
    :param src_operations: The source operations object.
    :type src_operations:
        :py:class:`delphixpy.v1_11_3.web.objects.LinkedSourceOperations`
        or :py:class:`delphixpy.v1_11_3.web.objects.VirtualSourceOperations`
    :param valid_op: A valid hook operation for an MSSql linked or
        virtual source.
    :type valid_op: ``str``
    :param hook_type: The name of the hook type to which the hooks needs to
        be migrated.
    :type hook_type: ``str``
    :returns: A bool value to indicate whether the given database source
        object has hooks in the given operation type(`valid_op`).
    :rtype: ``bool``
    """
    hook_present = False
    if hasattr(database_src.operations, valid_op):

        hooks_in_valid_op = getattr(database_src.operations, valid_op)
        for hook in hooks_in_valid_op:
            hook.type = hook_type
            hook_present = True

        setattr(src_operations, valid_op, hooks_in_valid_op)
    else:
        # Empty list indicates that the given database source object
        # doesn't have any hooks for the operation type- `valid_op`.
        setattr(src_operations, valid_op, [])

    return hook_present


def disable_feature_flag(engine_obj, dlpx_host_address):
    """
    The function disables the 'INSTALLEDPOWERSHELL' feature flag on the
    given engine.

    :param engine_obj: The Delphix engine object.
    :type engine_obj:
        :py:class:`delphixpy.v1_11_3.delphix_engine.DelphixEngine`
    :param dlpx_host_address: The host address of the engine.
    :type dlpx_host_address: ``str``
    """

    ps2_feature_flag = vo.FeatureFlagParameters()
    ps2_feature_flag.name = FEATURE_FLAG_INSTALLEDPOWERSHELL
    print(
        "\nDisabling the feature flag- '{flag}'"
        " on the engine- [{engine}] ...\n".format(
            flag=FEATURE_FLAG_INSTALLEDPOWERSHELL, engine=dlpx_host_address
        )
    )
    delphixpy.v1_11_3.web.system.disable_feature_flag(
        engine_obj, ps2_feature_flag
    )


if __name__ == "__main__":
    main()