# Copyright (c) 2016, Almar Klein

"""
This module implements functionality to obtain registered
Python interpreters and to register a Python interpreter in the Windows
registry.

"""

import sys
import os

try:
    import winreg
except ImportError:
    winreg = None

PYTHON_KEY = "SOFTWARE\\Python\\PythonCore"
PYTHON_KEY_WOW64 = "SOFTWARE\\Wow6432Node\\Python\\PythonCore"
INSTALL_KEY = "InstallPath"
PATH_KEY = "PythonPath"


class PythonInReg:
    """Class to represent a specific version of the Python interpreter
    registered (or being registered in the registry).
    This is a helper class for the functions defined in this module; it
    should not be instantiated directly.
    """

    USER_ONE = 1
    USER_ALL = 2

    def __init__(self, user, version, wow64=False):
        self._user = user
        self._key = (wow64 and PYTHON_KEY_WOW64 or PYTHON_KEY) + "\\" + version

    def __repr__(self):
        userstr = [None, "USER_ONE", "USER_ALL"][self._user]
        installPath = self.installPath()
        reg = self._reg()
        if not reg:
            return "<PythonInReg {} at {} (unregistered)>".format(self.version(), userstr)
        elif installPath:
            return '<PythonInReg {} at {} in "{}">'.format(
                self.version(),
                userstr,
                installPath,
            )
        else:
            return "<PythonInReg {} at {}>".format(self.version(), userstr)

    def _root(self):
        if self._user == PythonInReg.USER_ONE:
            return winreg.HKEY_CURRENT_USER
        else:
            return winreg.HKEY_LOCAL_MACHINE

    def _reg(self):
        # Get key for this version
        try:
            return winreg.OpenKey(self._root(), self._key, 0, winreg.KEY_READ)
        except Exception:
            return None

    def create(self):
        """Create key. If already exists, does nothing."""

        # Get key for this version
        reg = self._reg()

        if reg:
            winreg.CloseKey(reg)
            # print("Unable to create Python version {}: already exists.".format(self.version()))

        else:
            # Try to create
            try:
                reg = winreg.CreateKey(self._root(), self._key)
                winreg.CloseKey(reg)
            except Exception:
                raise RuntimeError(
                    "Unable to create python version {}.".format(self.version())
                )
            print("Created {}.".format(self))

    def delete(self):
        # Get key for this version
        reg = self._reg()
        if not reg:
            print("Unable to delete Python version {}: does not exist.".format(self.version()))

        # Delete attributes
        try:
            winreg.DeleteKey(reg, INSTALL_KEY)
        except Exception:
            pass
        try:
            winreg.DeleteKey(reg, PATH_KEY)
        except Exception:
            pass

        # Delete main key for this version, or show warning
        try:
            winreg.DeleteKey(self._root(), self._key)
        except Exception:
            print("Could not delete {}.".format(self))
            return
        print("Deleted {}.".format(self))

    def setInstallPath(self, installPath):
        # Get key for this version
        reg = self._reg()
        if not reg:
            raise RuntimeError(
                "Could not set installPath for version {}: version does not exist.".format(self.version())
            )

        # Set value or raise error
        try:
            winreg.SetValue(reg, INSTALL_KEY, winreg.REG_SZ, installPath)
            winreg.CloseKey(reg)
        except Exception:
            winreg.CloseKey(reg)
            raise RuntimeError("Could not set installPath for {}.".format(self))

    def installPath(self):
        # Get key for this version
        reg = self._reg()
        if not reg:
            return None

        # Get value or return None
        try:
            installPath = winreg.QueryValue(reg, INSTALL_KEY)
            winreg.CloseKey(reg)
            return installPath
        except Exception:
            winreg.CloseKey(reg)
            return None

    def setPythonPath(self, pythonPath):
        # Get key for this version
        reg = self._reg()
        if not reg:
            raise RuntimeError(
                "Could not set pythonPath for version {}: version does not exist.".format(self.version())
            )

        # Set value or raise error
        try:
            winreg.SetValue(reg, PATH_KEY, winreg.REG_SZ, pythonPath)
            winreg.CloseKey(reg)
        except Exception:
            winreg.CloseKey(reg)
            raise RuntimeError("Could not set pythonPath for {}.".format(self))

    def pythonPath(self):
        # Get key for this version
        reg = self._reg()
        if not reg:
            return None

        # Get value or return None
        try:
            pythonPath = winreg.QueryValue(reg, PATH_KEY)
            winreg.CloseKey(reg)
            return pythonPath
        except Exception:
            winreg.CloseKey(reg)
            return None

    def version(self):
        """Get the Python version."""
        return self._key[-3:]


def get_interpreters_in_reg():
    """get_interpreters_in_reg()
    Get a list of PythonInReg instances: one for each interpreter
    in the registry. This function checks both LOCAL_MACHINE and CURRENT_USER.
    """
    versions = []
    for user in [1, 2]:
        for wow64 in [False, True]:
            versions.extend(_get_interpreter_in_reg(user, wow64))
    return versions


def _get_interpreter_in_reg(user, wow64=False):
    # Get base key
    if user == PythonInReg.USER_ONE:
        HKEY = winreg.HKEY_CURRENT_USER
    else:
        HKEY = winreg.HKEY_LOCAL_MACHINE

    # Get Python key
    if wow64:
        PYKEY = PYTHON_KEY_WOW64
    else:
        PYKEY = PYTHON_KEY

    # Try to open Python key
    try:
        reg = winreg.OpenKey(HKEY, PYKEY, 0, winreg.KEY_READ)
    except Exception:
        return []

    # Get info about subkeys
    nsub, nval, modified = winreg.QueryInfoKey(reg)

    # Query all
    versions = []
    for i in range(nsub):
        # Get name and subkey
        version = winreg.EnumKey(reg, i)
        versions.append(PythonInReg(user, version, wow64))

    # Done
    winreg.CloseKey(reg)
    return versions


def register_interpreter(version=None, installPath=None, user=None, wow64=False):
    """register_interpreter(version=None, installPath=None, user=None, wow64=False)
    Register a certain Python version. If version and installPath
    are not given, the current Python process is registered.
    if user is not given, tries LOCAL_MACHINE first but uses CURRENT_USER
    if that fails.
    """
    if version is None:
        version = sys.version[:3]
    if installPath is None:
        installPath = sys.prefix

    # Get existing versions
    existingVersions = get_interpreters_in_reg()

    # Determine what users to try
    if user is None:
        users = [2, 1]
    else:
        users = [user]

    success = False

    for user in users:
        # Create new PythonInReg instance
        v = PythonInReg(user, version, wow64)

        # Check if already exists
        ok = True
        for ev in existingVersions:
            if ev._key != v._key or ev._user != v._user:
                continue  # Different key; no problem
            if (not ev.installPath()) or (not os.path.isdir(ev.installPath())):
                continue  # Key the same, but existing entry is invalid
            if ev.installPath() == installPath:
                # Exactly the same, no action required, return now!
                return ev
            # Ok, there's a problem
            ok = False
            print(
                'Warning: version {} is already installed in "{}".'.format(version, ev.installPath())
            )
        if not ok:
            continue

        # Try to create the key
        try:
            v.create()
            v.setInstallPath(installPath)
            success = True
            break
        except RuntimeError:
            continue

    if success:
        return v
    else:
        raise RuntimeError(
            "Could not register Python version {} at {}.".format(version, installPath)
        )


if __name__ == "__main__":
    for v in get_interpreters_in_reg():
        print(v)
