# Copyright (c) 2018-2019 VMware, Inc.  All rights reserved.
# -- VMware Confidential

"""VCDB VMODL In-Place (B2B) Upgrade PyVMOMI abstraction

This module defines mechanisms to import more than one pyVmomi in the current
process. This is necessary for testing purposes as well as for live runs.
"""

import os
import sys


class _TempSysPath(object):
    """Alter system paths so that a pyVmomi path is temporary available

    This generic utility enables multiple pyVmomi instances.
    """

    def __init__(self, path):
        self.path = path

    def __enter__(self):
        sys.path.insert(0, self.path)

    def __exit__(self, exc_type, exc_value, traceback):
        sys.path.remove(self.path)


def _ImportPyVmomi():
    """Import a pyVmomi instance depending on the current PYTHONPATH.

    Also, import the VSAN-related modules that are not imported by default.
    """

    # pylint: disable=import-error
    # pylint: disable=bare-except
    import pyVmomi
    import pyVmomi.VmomiSupport

    try:
        # Add dynamically the VSAN mgmt types. @see PR 1904479, PR 1790123
        import pyVmomi.vsanmgmtObjects
    except:
        # Can't find them. Hopefully, they don't have anything in VCDB.
        pass

    for m in [m for m in sys.modules if m.startswith("pyVmomi")]:
        del sys.modules[m]

    return pyVmomi


def MakePyVmomiImporter(path):
    """Creates a custom function to import pyVmomi from path

    This is necessary for being able to mock the In-Place upgrade infra.
    """

    def result():
        with _TempSysPath(path):
            return _ImportPyVmomi()

    return result


def ImportSrcPyVmomi():
    """Import the default source pyVmomi

    It is packaged on the prePatch step of the In-Place upgrade.
    Currently, it is copied next to __file__ because it was easiest at the
    time of writing.
    """

    root = os.path.dirname(os.path.realpath(__file__))
    importer = MakePyVmomiImporter(os.path.join(root, "oldPyVmomi"))
    return importer()


def ImportDstPyVmomi():
    """Import the default destination pyVmomi

    It is packaged in the patch metadata scripts of the destination build.
    Currently, it is extraceted next to __file__ because it was easiest at the
    time of writing.
    """

    root = os.path.dirname(os.path.realpath(__file__))
    importer = MakePyVmomiImporter(os.path.join(root, "currentPyVpx"))
    return importer()
