# Copyright 2016-2020 Swiss National Supercomputing Centre (CSCS/ETH Zurich)
# ReFrame Project Developers. See the top-level LICENSE file for details.
#
# SPDX-License-Identifier: BSD-3-Clause

#
# Utilities for manipulating the modules subsystem
#

import abc
import os
import re
from collections import OrderedDict

import reframe.core.fields as fields
import reframe.utility.osext as osext
import reframe.utility.typecheck as types
from reframe.core.exceptions import (ConfigError, EnvironError,
                                     SpawnedProcessError)
from reframe.utility import OrderedSet


class Module:
    '''Module wrapper.

    This class represents internally a module. Concrete module system
    implementation should deal only with that.
    '''

    def __init__(self, name):
        if not isinstance(name, str):
            raise TypeError('module name not a string')

        name = name.strip()
        if not name:
            raise ValueError('module name cannot be empty')

        try:
            self._name, self._version = name.split('/', maxsplit=1)
        except ValueError:
            self._name, self._version = name, None

    @property
    def name(self):
        return self._name

    @property
    def version(self):
        return self._version

    @property
    def fullname(self):
        if self.version is not None:
            return '/'.join((self.name, self.version))
        else:
            return self.name

    def __hash__(self):
        # Here we hash only over the name of the module, because foo/1.2 and
        # simply foo compare equal. In case of hash conflicts (e.g., foo/1.2
        # and foo/1.3), the equality operator will resolve it.
        return hash(self.name)

    def __eq__(self, other):
        if not isinstance(other, type(self)):
            return NotImplemented

        if not self.version or not other.version:
            return self.name == other.name
        else:
            return self.name == other.name and self.version == other.version

    def __repr__(self):
        return '%s(%s)' % (type(self).__name__, self.fullname)

    def __str__(self):
        return self.fullname


class ModulesSystem:
    '''A modules system.'''

    module_map = fields.TypedField('module_map',
                                   types.Dict[str, types.List[str]])

    @classmethod
    def create(cls, modules_kind=None):
        if modules_kind is None or modules_kind == 'nomod':
            return ModulesSystem(NoModImpl())
        elif modules_kind == 'tmod31':
            return ModulesSystem(TMod31Impl())
        elif modules_kind == 'tmod':
            return ModulesSystem(TModImpl())
        elif modules_kind == 'tmod32':
            return ModulesSystem(TModImpl())
        elif modules_kind == 'tmod4':
            return ModulesSystem(TMod4Impl())
        elif modules_kind == 'lmod':
            return ModulesSystem(LModImpl())
        else:
            raise ConfigError('unknown module system: %s' % modules_kind)

    def __init__(self, backend):
        self._backend = backend
        self.module_map = {}

    def resolve_module(self, name):
        '''Resolve module ``name`` in the registered module map.

        :returns: the list of real modules names pointed to by ``name``.
        :raises: :class:`reframe.core.exceptions.ConfigError` if the mapping
            contains a cycle.

        :meta private:
        '''
        ret = OrderedSet()
        visited = set()
        unvisited = [(name, None)]
        path = []
        while unvisited:
            node, parent = unvisited.pop()
            # Adjust the path
            while path and path[-1] != parent:
                path.pop()

            # Handle modules mappings with self loops
            if node == parent:
                ret.add(node)
                continue

            try:
                # We insert the adjacent nodes in reverse order, so as to
                # preserve the DFS access order
                adjacent = reversed(self.module_map[node])
            except KeyError:
                # We have reached a terminal node
                ret.add(node)
            else:
                path.append(node)
                for m in adjacent:
                    if m in path and m != node:
                        raise EnvironError('module cyclic dependency: ' +
                                           '->'.join(path + [m]))
                    if m not in visited:
                        unvisited.append((m, node))

            visited.add(node)

        return list(ret)

    @property
    def backend(self):
        return(self._backend)

    def available_modules(self, substr=None):
        '''Return a list of available modules that contain ``substr`` in their
        name.

        :rtype: List[str]
        '''
        return [str(m) for m in self._backend.available_modules(substr or '')]

    def loaded_modules(self):
        '''Return a list of loaded modules.

        :rtype: List[str]
        '''
        return [str(m) for m in self._backend.loaded_modules()]

    def conflicted_modules(self, name):
        '''Return the list of the modules conflicting with module ``name``.

        If module ``name`` resolves to multiple real modules, then the returned
        list will be the concatenation of the conflict lists of all the real
        modules.

        :rtype: List[str]
        '''
        ret = []
        for m in self.resolve_module(name):
            ret += self._conflicted_modules(m)

        return ret

    def _conflicted_modules(self, name):
        return [str(m) for m in self._backend.conflicted_modules(Module(name))]

    def load_module(self, name, force=False):
        '''Load the module ``name``.

        If ``force`` is set, forces the loading, unloading first any
        conflicting modules currently loaded. If module ``name`` refers to
        multiple real modules, all of the target modules will be loaded.

        :returns: the list of unloaded modules as strings.
        :rtype: List[str]
        '''
        ret = []
        for m in self.resolve_module(name):
            ret += self._load_module(m, force)

        return ret

    def _load_module(self, name, force=False):
        module = Module(name)
        loaded_modules = self._backend.loaded_modules()
        if module in loaded_modules:
            # Do not try to load the module if it is already present
            return []

        # Get the list of the modules that need to be unloaded
        unload_list = set()
        if force:
            conflict_list = self._backend.conflicted_modules(module)
            unload_list = set(loaded_modules) & set(conflict_list)

        for m in unload_list:
            self._backend.unload_module(m)

        self._backend.load_module(module)
        return [str(m) for m in unload_list]

    def unload_module(self, name):
        '''Unload module ``name``.

        If module ``name`` refers to multiple real modules, all the referred to
        modules will be unloaded in reverse order.
        '''
        for m in reversed(self.resolve_module(name)):
            self._unload_module(m)

    def _unload_module(self, name):
        self._backend.unload_module(Module(name))

    def is_module_loaded(self, name):
        '''Check if module ``name`` is loaded.

        If module ``name`` refers to multiple real modules, this method will
        return :class:`True` only if all the referees are loaded.
        '''
        return all(self._is_module_loaded(m)
                   for m in self.resolve_module(name))

    def _is_module_loaded(self, name):
        return self._backend.is_module_loaded(Module(name))

    def load_mapping(self, mapping):
        '''Update the internal module mappings using a single mapping.

        :arg mapping: a string specifying the module mapping.
            Example syntax: ``'m0: m1 m2'``.

        :meta private:
        '''
        key, *rest = mapping.split(':')
        if len(rest) != 1:
            raise ConfigError('invalid mapping syntax: %s' % mapping)

        key = key.strip()
        values = rest[0].split()
        if not key:
            raise ConfigError('no key found in mapping: %s' % mapping)

        if not values:
            raise ConfigError('no mapping defined for module: %s' % key)

        self.module_map[key] = list(OrderedDict.fromkeys(values))

    def load_mapping_from_file(self, filename):
        '''Update the internal module mappings from mappings read from file.

        :meta private:
        '''
        with open(filename) as fp:
            for lineno, line in enumerate(fp, start=1):
                line = line.strip().split('#')[0]
                if not line:
                    continue

                try:
                    self.load_mapping(line)
                except ConfigError as e:
                    raise ConfigError('%s:%s' % (filename, lineno)) from e

    @property
    def name(self):
        '''The name of this module system.'''
        return self._backend.name()

    @property
    def version(self):
        '''The version of this module system.'''
        return self._backend.version()

    def unload_all(self):
        '''Unload all loaded modules.'''
        return self._backend.unload_all()

    @property
    def searchpath(self):
        '''The module system search path as a list of directories.'''
        return self._backend.searchpath()

    def searchpath_add(self, *dirs):
        '''Add ``dirs`` to the module system search path.'''
        return self._backend.searchpath_add(*dirs)

    def searchpath_remove(self, *dirs):
        '''Remove ``dirs`` from the module system search path.'''
        return self._backend.searchpath_remove(*dirs)

    def emit_load_commands(self, name):
        '''Return the appropriate shell command for loading module ``name``.

        :rtype: List[str]
        '''
        return [self._backend.emit_load_instr(Module(name))
                for name in self.resolve_module(name)]

    def emit_unload_commands(self, name):
        '''Return the appropriate shell command for unloading module
        ``name``.

        :rtype: List[str]
        '''
        return [self._backend.emit_unload_instr(Module(name))
                for name in reversed(self.resolve_module(name))]

    def __str__(self):
        return str(self._backend)


class ModulesSystemImpl(abc.ABC):
    '''Abstract base class for module systems.'''

    @abc.abstractmethod
    def available_modules(self, substr):
        '''Return a list of available modules, whose name contains ``substr``.

        This method returns a list of Module instances.
        '''

    @abc.abstractmethod
    def loaded_modules(self):
        '''Return a list of loaded modules.

        This method returns a list of Module instances.
        '''

    @abc.abstractmethod
    def conflicted_modules(self, module):
        '''Return the list of conflicted modules.

        This method returns a list of Module instances.
        '''

    @abc.abstractmethod
    def load_module(self, module):
        '''Load the module ``name``.

        If ``force`` is set, forces the loading,
        unloading first any conflicting modules currently loaded.

        Returns the unloaded modules as a list of module instances.'''

    @abc.abstractmethod
    def unload_module(self, module):
        '''Unload module ``module``.'''

    @abc.abstractmethod
    def is_module_loaded(self, module):
        '''Check presence of module ``module``.'''

    @abc.abstractmethod
    def name(self):
        '''Return the name of this module system.'''

    @abc.abstractmethod
    def version(self):
        '''Return the version of this module system.'''

    @abc.abstractmethod
    def unload_all(self):
        '''Unload all loaded modules.'''

    @abc.abstractmethod
    def searchpath(self):
        '''The module system search path as a list of directories.'''

    @abc.abstractmethod
    def searchpath_add(self, *dirs):
        '''Add ``dirs`` to the module system search path.'''

    @abc.abstractmethod
    def searchpath_remove(self, *dirs):
        '''Remove ``dirs`` from the module system search path.'''

    @abc.abstractmethod
    def emit_load_instr(self, module):
        '''Emit the instruction that loads module.'''

    @abc.abstractmethod
    def emit_unload_instr(self, module):
        '''Emit the instruction that unloads module.'''

    def __repr__(self):
        return type(self).__name__ + '()'

    def __str__(self):
        return self.name() + ' ' + self.version()


class TModImpl(ModulesSystemImpl):
    '''Base class for TMod Module system (Tcl).'''

    MIN_VERSION = (3, 2)

    def __init__(self):
        # Try to figure out if we are indeed using the TCL version
        try:
            completed = osext.run_command('modulecmd -V')
        except OSError as e:
            raise ConfigError(
                'could not find a sane TMod installation') from e

        version_match = re.search(r'^VERSION=(\S+)', completed.stdout,
                                  re.MULTILINE)
        tcl_version_match = re.search(r'^TCL_VERSION=(\S+)', completed.stdout,
                                      re.MULTILINE)

        if version_match is None or tcl_version_match is None:
            raise ConfigError('could not find a sane TMod installation')

        version = version_match.group(1)
        try:
            ver_major, ver_minor = [int(v) for v in version.split('.')[:2]]
        except ValueError:
            raise ConfigError(
                'could not parse TMod version string: ' + version) from None

        if (ver_major, ver_minor) < self.MIN_VERSION:
            raise ConfigError(
                'unsupported TMod version: %s (required >= %s)' %
                (version, self.MIN_VERSION))

        self._version = version
        self._command = 'modulecmd python'
        try:
            # Try the Python bindings now
            completed = osext.run_command(self._command)
        except OSError as e:
            raise ConfigError(
                'could not get the Python bindings for TMod: ' % e) from e

        if re.search(r'Unknown shell type', completed.stderr):
            raise ConfigError(
                'Python is not supported by this TMod installation')

    def name(self):
        return 'tmod'

    def version(self):
        return self._version

    def _run_module_command(self, *args, msg=None):
        command = ' '.join([self._command, *args])
        try:
            completed = osext.run_command(command, check=True)
        except SpawnedProcessError as e:
            raise EnvironError(msg) from e

        if self._module_command_failed(completed):
            err = SpawnedProcessError(command,
                                      completed.stdout,
                                      completed.stderr,
                                      completed.returncode)
            raise EnvironError(msg) from err

        return completed

    def _module_command_failed(self, completed):
        return re.search(r'ERROR', completed.stderr) is not None

    def _exec_module_command(self, *args, msg=None):
        completed = self._run_module_command(*args, msg=msg)
        exec(completed.stdout)

    def available_modules(self, substr):
        completed = self._run_module_command(
            'avail', '-t', substr, msg='could not retrieve available modules'
        )
        ret = []
        for line in completed.stderr.split('\n'):
            if not line or line[-1] == ':':
                # Ignore empty lines and path entries
                continue

            module = re.sub(r'\(default\)', '', line)
            ret.append(Module(module))

        return ret

    def loaded_modules(self):
        try:
            # LOADEDMODULES may be defined but empty
            return [Module(m)
                    for m in os.environ['LOADEDMODULES'].split(':') if m]
        except KeyError:
            return []

    def conflicted_modules(self, module):
        completed = self._run_module_command(
            'show', str(module), msg="could not show module '%s'" % module)
        return [Module(m.group(1))
                for m in re.finditer(r'^conflict\s+(\S+)',
                                     completed.stderr, re.MULTILINE)]

    def is_module_loaded(self, module):
        return module in self.loaded_modules()

    def load_module(self, module):
        self._exec_module_command(
            'load', str(module),
            msg="could not load module '%s' correctly" % module)

    def unload_module(self, module):
        self._exec_module_command(
            'unload', str(module),
            msg="could not unload module '%s' correctly" % module)

    def unload_all(self):
        self._exec_module_command('purge')

    def searchpath(self):
        path = os.getenv('MODULEPATH', '')
        return path.split(':')

    def searchpath_add(self, *dirs):
        self._exec_module_command('use', *dirs)

    def searchpath_remove(self, *dirs):
        self._exec_module_command('unuse', *dirs)

    def emit_load_instr(self, module):
        return 'module load %s' % module

    def emit_unload_instr(self, module):
        return 'module unload %s' % module


class TMod31Impl(TModImpl):
    '''Module system for TMod (Tcl).'''

    MIN_VERSION = (3, 1)

    def __init__(self):
        # Try to figure out if we are indeed using the TCL version
        try:
            modulecmd = os.getenv('MODULESHOME')
            modulecmd = os.path.join(modulecmd, 'modulecmd.tcl')
            completed = osext.run_command(modulecmd)
        except OSError as e:
            raise ConfigError(
                'could not find a sane TMod31 installation: %s' % e) from e

        version_match = re.search(r'Release Tcl (\S+)', completed.stderr,
                                  re.MULTILINE)
        tcl_version_match = version_match

        if version_match is None or tcl_version_match is None:
            raise ConfigError('could not find a sane TMod31 installation')

        version = version_match.group(1)
        try:
            ver_major, ver_minor = [int(v) for v in version.split('.')[:2]]
        except ValueError:
            raise ConfigError(
                'could not parse TMod31 version string: ' + version) from None

        if (ver_major, ver_minor) < self.MIN_VERSION:
            raise ConfigError(
                'unsupported TMod version: %s (required >= %s)' %
                (version, self.MIN_VERSION))

        self._version = version
        self._command = '%s python' % modulecmd

        try:
            # Try the Python bindings now
            completed = osext.run_command(self._command)
        except OSError as e:
            raise ConfigError(
                'could not get the Python bindings for TMod31: ' % e) from e

        if re.search(r'Unknown shell type', completed.stderr):
            raise ConfigError(
                'Python is not supported by this TMod installation')

    def name(self):
        return 'tmod31'

    def _exec_module_command(self, *args, msg=None):
        completed = self._run_module_command(*args, msg=msg)
        exec_match = re.search(r'^exec\s\'', completed.stdout)
        if exec_match is None:
            raise ConfigError('could not use the python bindings')
        else:
            cmd = completed.stdout
            exec_match = re.search(r'^exec\s\'(\S+)\'', cmd,
                                   re.MULTILINE)
            if exec_match is None:
                raise ConfigError('could not use the python bindings')
            with open(exec_match.group(1), 'r') as content_file:
                cmd = content_file.read()

        exec(cmd)


class TMod4Impl(TModImpl):
    '''Module system for TMod 4.'''

    MIN_VERSION = (4, 1)

    def __init__(self):
        self._command = 'modulecmd python'
        try:
            completed = osext.run_command(self._command + ' -V', check=True)
        except OSError as e:
            raise ConfigError(
                'could not find a sane TMod4 installation') from e
        except SpawnedProcessError as e:
            raise ConfigError(
                'could not get the Python bindings for TMod4') from e

        version_match = re.match(r'^Modules Release (\S+)\s+',
                                 completed.stderr)
        if not version_match:
            raise ConfigError('could not retrieve the TMod4 version')

        version = version_match.group(1)
        try:
            ver_major, ver_minor = [int(v) for v in version.split('.')[:2]]
        except ValueError:
            raise ConfigError(
                'could not parse TMod4 version string: ' + version) from None

        if (ver_major, ver_minor) < self.MIN_VERSION:
            raise ConfigError(
                'unsupported TMod4 version: %s (required >= %s)' %
                (version, self.MIN_VERSION))

        self._version = version

    def name(self):
        return 'tmod4'

    def _exec_module_command(self, *args, msg=None):
        command = ' '.join([self._command, *args])
        completed = osext.run_command(command, check=True)
        namespace = {}
        exec(completed.stdout, {}, namespace)
        if not namespace['_mlstatus']:
            # _mlstatus is set by the TMod4 Python bindings
            if msg is None:
                msg = 'modules system command failed: '
                if isinstance(completed.args, str):
                    msg += completed.args
                else:
                    msg += ' '.join(completed.args)

            raise EnvironError(msg)


class LModImpl(TModImpl):
    '''Module system for Lmod (Tcl/Lua).'''

    def __init__(self):
        # Try to figure out if we are indeed using LMOD
        lmod_cmd = os.getenv('LMOD_CMD')
        if lmod_cmd is None:
            raise ConfigError('could not find a sane Lmod installation: '
                              'environment variable LMOD_CMD is not defined')

        try:
            completed = osext.run_command('%s --version' % lmod_cmd)
        except OSError as e:
            raise ConfigError(
                'could not find a sane Lmod installation: %s' % e)

        version_match = re.search(r'.*Version\s*(\S+)', completed.stderr,
                                  re.MULTILINE)
        if version_match is None:
            raise ConfigError('could not retrieve Lmod version')

        self._version = version_match.group(1)
        self._command = '%s python ' % lmod_cmd
        try:
            # Try the Python bindings now
            completed = osext.run_command(self._command)
        except OSError as e:
            raise ConfigError(
                'could not get the Python bindings for Lmod: ' % e)

        if re.search(r'Unknown shell type', completed.stderr):
            raise ConfigError('Python is not supported by '
                              'this Lmod installation')

    def name(self):
        return 'lmod'

    def _module_command_failed(self, completed):
        return completed.stdout.strip() == 'false'

    def available_modules(self, substr):
        completed = self._run_module_command(
            '-t', 'avail', substr, msg='could not retrieve available modules'
        )
        ret = []
        for line in completed.stderr.split('\n'):
            if not line or line[-1] == ':':
                # Ignore empty lines and path entries
                continue

            module = re.sub(r'\(\S+\)', '', line)
            ret.append(Module(module))

        return ret

    def conflicted_modules(self, module):
        completed = self._run_module_command(
            'show', str(module), msg="could not show module '%s'" % module)

        # Lmod accepts both Lua and and Tcl syntax
        # The following test allows incorrect syntax, e.g., `conflict
        # ('package"(`, but we expect this to be caught by the Lmod framework
        # in earlier stages.
        ret = []
        for m in re.finditer(r'conflict\s*(\S+)', completed.stderr):
            conflict_arg = m.group(1)
            if conflict_arg.startswith('('):
                # Lua syntax
                ret.append(Module(conflict_arg.strip('\'"()')))
            else:
                # Tmod syntax
                ret.append(Module(conflict_arg))

        return ret

    def unload_all(self):
        # Currently, we don't take any provision for sticky modules in Lmod, so
        # we forcefully unload everything.
        self._exec_module_command('--force', 'purge')


class NoModImpl(ModulesSystemImpl):
    '''A convenience class that implements a no-op a modules system.'''

    def available_modules(self, substr):
        return []

    def loaded_modules(self):
        return []

    def conflicted_modules(self, module):
        return []

    def load_module(self, module):
        pass

    def unload_module(self, module):
        pass

    def is_module_loaded(self, module):
        #
        # Always return `True`, since this pseudo modules system effectively
        # assumes that everything needed is loaded.
        #
        return True

    def name(self):
        return 'nomod'

    def version(self):
        return '1.0'

    def unload_all(self):
        pass

    def searchpath(self):
        return []

    def searchpath_add(self, *dirs):
        pass

    def searchpath_remove(self, *dirs):
        pass

    def emit_load_instr(self, module):
        return ''

    def emit_unload_instr(self, module):
        return ''
