# *****************************************************************************
#
#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.
#
#   See NOTICE file for details.
#
# *****************************************************************************
#   Copyright 2013 Thomas Calmant

import os
import sys

__all__ = ['getDefaultJVMPath',
           'get_default_jvm_path',
           'JVMNotFoundException',
           'JVMNotSupportedException']

from typing import Sequence, Tuple


class JVMNotFoundException(ValueError):
    """Exception raised when no JVM was found in the search path.

    This exception is raised when the all of the places searched did not
    contain a JVM. The locations searched depend on the machine architecture.
    To avoid this exception specify the JAVA_HOME environment variable as a
    valid jre or jdk root directory.
    """


class JVMNotSupportedException(ValueError):
    """Exception raised when the JVM is not supported.

    This exception is raised after a search found a valid Java home directory
    was found, but the JVM shared library found is not supported. Typically
    this occurs when the JVM does not match the architecture of Python
    32 vs 64 bit, or the JVM is older than the version used to compile
    JPype.
    """


def getDefaultJVMPath() -> str:
    """Retrieves the path to the default or first found JVM library.

    Returns:
      The path to the JVM shared library file

    Raises:
      JVMNotFoundException: If there was no JVM found in the search path.
      JVMNotSupportedException: If the JVM was found was not compatible with
        Python due to cpu architecture.

    """
    if sys.platform == "win32":
        finder = WindowsJVMFinder()
    elif sys.platform == "darwin":
        finder = DarwinJVMFinder()
    else:
        finder = LinuxJVMFinder()
    return finder.get_jvm_path()


get_default_jvm_path = getDefaultJVMPath


class JVMFinder:
    """JVM library finder base class."""
    # Library file name
    _libfile: str = "libjvm.so"

    # Predefined locations
    _locations: Tuple[str, ...] = ("/usr/lib/jvm", "/usr/java")

    def __init__(self):
        # Search methods
        self._methods = (self._get_from_java_home,
                         self._get_from_known_locations)

    def find_libjvm(self, java_home):
        """Recursively looks for the given file.

        Parameters:
            java_home(str): A Java home folder

        Returns:
            The first found file path, or None
        """
        non_supported_jvm = ('cacao', 'jamvm')
        found_non_supported_jvm = False

        # Look for the file
        for root, _, names in sorted(os.walk(java_home), key=lambda t: len(t[0].split(os.sep))):
            if self._libfile in names:
                # Found it, but check for non supported jvms
                candidate = os.path.split(root)[1]
                if candidate in non_supported_jvm:
                    found_non_supported_jvm = True
                    continue  # maybe we will find another one?
                return os.path.join(root, self._libfile)

        if found_non_supported_jvm:
            raise JVMNotSupportedException("Sorry '{0}' is known to be "
                                           "broken. Please ensure your "
                                           "JAVA_HOME contains at least "
                                           "another JVM implementation "
                                           "(eg. server)"
                                           .format(candidate))
        # File not found
        raise JVMNotFoundException("Sorry no JVM could be found. "
                                   "Please ensure your JAVA_HOME "
                                   "environment variable is pointing "
                                   "to correct installation.")

    @staticmethod
    def find_possible_homes(parents):
        """
        Generator that looks for the first-level children folders that could be
        Java installations, according to their name

        Parameters:
            parents (str[]): A list of parent directories

        Returns:
            A list of the possible JVM installation folders
        """
        homes = []
        java_names = ('jre', 'jdk', 'java')

        for parent in parents:
            # Fast exit if folder does not exist
            if not os.path.exists(parent):
                continue

            for childname in sorted(os.listdir(parent)):
                # Compute the real path
                path = os.path.realpath(os.path.join(parent, childname))
                if path in homes or not os.path.isdir(path):
                    # Already known path, or not a directory -> ignore
                    continue

                # Check if the path seems OK
                real_name = os.path.basename(path).lower()
                for java_name in java_names:
                    if java_name in real_name:
                        # Correct JVM folder name
                        homes.append(path)
                        yield path
                        break

    def check(self, jvm):
        """
        Check if the jvm is valid for this architecture.

        This method should be overriden for each architecture.

        Raises:
            JVMNotSupportedException: If the jvm is not supported.
        """
        pass

    def get_jvm_path(self):
        """
        Retrieves the path to the default or first found JVM library

        Returns:
            The path to the JVM shared library file

        Raises:
            ValueError: No JVM library found or No Support JVM found
        """
        jvm_notsupport_ext = None
        for method in self._methods:
            try:
                jvm = method()

                # If found check the architecture
                if jvm:
                    self.check(jvm)
            except NotImplementedError:
                # Ignore missing implementations
                pass
            except JVMNotFoundException:
                # Ignore not successful methods
                pass
            except JVMNotSupportedException as e:
                jvm_notsupport_ext = e

            else:
                if jvm is not None:
                    return jvm

        if jvm_notsupport_ext is not None:
            raise jvm_notsupport_ext
        raise JVMNotFoundException("No JVM shared library file ({0}) "
                                   "found. Try setting up the JAVA_HOME "
                                   "environment variable properly."
                                   .format(self._libfile))

    def _get_from_java_home(self):
        """
        Retrieves the Java library path according to the JAVA_HOME environment
        variable

        Returns:
            The path to the JVM library, or None
        """
        # Get the environment variable
        java_home = os.getenv("JAVA_HOME")
        if java_home and os.path.exists(java_home):
            # Get the real installation path
            java_home = os.path.realpath(java_home)

            if not os.path.exists(java_home):
                java_home = os.getenv("JAVA_HOME")

            # Look for the library file
            return self.find_libjvm(java_home)

    def _get_from_known_locations(self):
        """
        Retrieves the first existing Java library path in the predefined known
        locations

        Returns:
            The path to the JVM library, or None
        """
        for home in self.find_possible_homes(self._locations):
            jvm = self.find_libjvm(home)
            if jvm is not None:
                return jvm


class LinuxJVMFinder(JVMFinder):
    """Linux JVM library finder class."""

    # Java bin file
    _java = "/usr/bin/java"

    # Library file name
    _libfile = "libjvm.so"

    # Predefined locations
    _locations = ("/usr/lib/jvm", "/usr/java", "/opt/sun")

    def __init__(self):
        super().__init__()

        # Search methods
        self._methods = (self._get_from_java_home,
                         self._get_from_bin,
                         self._get_from_known_locations)

    def _get_from_bin(self):
        """
        Retrieves the Java library path according to the real installation of
        the java executable

        :return: The path to the JVM library, or None
        """
        # Find the real interpreter installation path
        java_bin = os.path.realpath(self._java)
        if os.path.exists(java_bin):
            # Get to the home directory
            java_home = os.path.abspath(os.path.join(os.path.dirname(java_bin),
                                                     '..'))

            # Look for the JVM library
            return self.find_libjvm(java_home)


class DarwinJVMFinder(LinuxJVMFinder):
    """
    Mac OS X JVM library finder class
    """
    # Library file name
    _libfile = "libjli.dylib"
    # Predefined locations
    _locations = ('/Library/Java/JavaVirtualMachines',)  # type: ignore

    def __init__(self):
        """
        Sets up members
        """
        super().__init__()

        self._methods = list(self._methods)
        self._methods.append(self._javahome_binary)

    def _javahome_binary(self):
        """
        for osx > 10.5 we have the nice util /usr/libexec/java_home available. Invoke it and
        return its output. It seems this tool has been removed in osx 10.9.
        """
        import platform
        import subprocess
        from packaging.version import Version

        current = Version(platform.mac_ver()[0][:4])
        # TODO: check if the java_home tool is still available and fix the version boundaries.
        if Version('10.6') <= current:  # < Version('10.9'):
            return subprocess.check_output(
                ['/usr/libexec/java_home']).strip()


def _checkJVMArch(jvmPath, maxsize=sys.maxsize):
    import struct
    IMAGE_FILE_MACHINE_I386 = 332
    IMAGE_FILE_MACHINE_IA64 = 512
    IMAGE_FILE_MACHINE_AMD64 = 34404

    is64 = maxsize > 2**32
    with open(jvmPath, "rb") as f:
        s = f.read(2)
        if s != b"MZ":
            raise JVMNotSupportedException("JVM not valid")
        f.seek(60)
        s = f.read(4)
        header_offset = struct.unpack("<L", s)[0]
        f.seek(header_offset + 4)
        s = f.read(2)
        machine = struct.unpack("<H", s)[0]

    if machine == IMAGE_FILE_MACHINE_I386:
        if is64:
            raise JVMNotSupportedException(
                "JVM mismatch, python is 64 bit and JVM is 32 bit.")
    elif machine == IMAGE_FILE_MACHINE_IA64 or machine == IMAGE_FILE_MACHINE_AMD64:
        if not is64:
            raise JVMNotSupportedException(
                "JVM mismatch, python is 32 bit and JVM is 64 bit.")
    else:
        raise JVMNotSupportedException("Unable to determine JVM Type")


class WindowsJVMFinder(JVMFinder):
    """
    Windows JVM library finder class
    """
    reg_keys = [r"SOFTWARE\JavaSoft\Java Runtime Environment",
                r"SOFTWARE\JavaSoft\JRE",
                ]
    # Library file name
    _libfile = "jvm.dll"

    def __init__(self):
        super().__init__()

        # Search methods
        self._methods = (self._get_from_java_home, self._get_from_registry)

    def check(self, jvm):
        _checkJVMArch(jvm)

    @staticmethod
    def _get_from_registry():
        """
        Retrieves the path to the default Java installation stored in the
        Windows registry

        :return: The path found in the registry, or None
        """
        try:
            import winreg
        except ImportError:
            return None

        for location in WindowsJVMFinder.reg_keys:
            try:
                jreKey = winreg.OpenKey(winreg.HKEY_LOCAL_MACHINE, location)
                cv = winreg.QueryValueEx(jreKey, "CurrentVersion")
                versionKey = winreg.OpenKey(jreKey, cv[0])
                winreg.CloseKey(jreKey)

                cv = winreg.QueryValueEx(versionKey, "RuntimeLib")
                winreg.CloseKey(versionKey)

                return cv[0]
            except OSError:
                pass
        return None