###############################################################################
# (c) Copyright 2018 CERN                                                     #
#                                                                             #
# This software is distributed under the terms of the GNU General Public      #
# Licence version 3 (GPL Version 3), copied verbatim in the file "COPYING".   #
#                                                                             #
# In applying this licence, CERN does not waive the privileges and immunities #
# granted to it by virtue of its status as an Intergovernmental Organization  #
# or submit itself to any jurisdiction.                                       #
###############################################################################
'''
@author: Marco Clemencic <marco.clemencic@cern.ch>
'''

import os
import sys
import traceback
import xenv

import LbPlatformUtils

from LbEnv.ProjectEnv.lookup import (
    getEnvXmlPath, findProject, findDataPackage, getLCGRelocation,
    getHepToolsInfo, NotFoundError, getProjectNameVersion, findLCGForExt,
    listVersions, listPlatforms, EXTERNAL_PROJECTS)
from LbEnv.ProjectEnv.version import (isValidVersion, expandVersionAlias,
                                      DEFAULT_VERSION, LCGInfoName)
from LbEnv import fixProjectCase

HOST_PLATFORM = LbPlatformUtils.dirac_platform()


def supported(binary_tag):
    'return True if the current host supports the given binary_tag'
    return LbPlatformUtils.can_run(HOST_PLATFORM,
                                   LbPlatformUtils.requires(binary_tag))


auto_override_projects = []


def decodePkg(s):
    '''
    Translate a string declaring a data package into a pair (package, version).

    The user can specify the data package they want to use in different ways
    and in all cases we need to return the package name with "hat" (also when
    the hat is implied) and the version pattern.

    >>> decodePkg('Hat/Package')
    ('Hat/Package', '*')
    >>> decodePkg('Package v7r*')
    ('Package', 'v7r*')
    >>> decodePkg('Package v1r0 Hat')
    ('Hat/Package', 'v1r0')
    >>> decodePkg('Hat/Package.v1r0')
    ('Hat/Package', 'v1r0')
    >>> decodePkg('Hat.Package.v1r0')
    ('Hat/Package', 'v1r0')
    >>> decodePkg('a b c d')
    Traceback (most recent call last):
       ...
    ValueError: too many words in package declaration 'a b c d'
    '''
    # split on spaces
    l = s.split()
    if len(l) == 2:
        pkg, vers = l
    elif len(l) == 3:
        pkg, vers, hat = l
        pkg = '{0}/{1}'.format(hat, pkg)
    elif len(l) == 1:
        if '.' in s:
            pkg, vers = s.rsplit('.', 1)
            pkg = pkg.replace('.', '/')
        else:
            pkg = s
            vers = '*'
    else:
        raise ValueError('too many words in package declaration %r' % s)

    return pkg, vers


def projectExtraPath(projroot):
    '''
    Return any extra search path required by the project at 'projroot'.
    '''
    from LbEnv.ProjectEnv.options import SearchPath, NightlyPathEntry
    extra_path = SearchPath([])
    # drop the 'InstallArea' part of the path
    while 'InstallArea' in projroot:
        projroot = os.path.dirname(projroot)

    def extractList(filename, varname):
        if os.path.exists(filename):
            data = {}
            exec open(filename).read() in data  # IGNORE:W0122
            # Get the list and convert it to strings
            return filter(str, data.get(varname, []))
        else:
            return []

    # check for the Python digested search path
    spFile = os.path.join(projroot, 'searchPath.py')
    if os.path.exists(spFile):
        data = {}
        exec open(spFile).read() in data  # IGNORE:W0122
        extra_path = data['path']

    # check for a requested nightly slot
    build_conf = os.path.join(projroot, 'build.conf')
    if os.path.exists(build_conf):
        vals = dict(l.strip().split('=', 1)
                    for l in map(str.strip, open(build_conf))
                    if l and not l.startswith('#'))
        slot = vals.get('nightly_slot')
        day = vals.get('nightly_day')
        base = vals.get('nightly_base') or os.environ.get('LHCBNIGHTLY', '')
        if slot and day and base:
            for p in extra_path:
                if isinstance(p, NightlyPathEntry):
                    p.base, p.slot, p.day = base, slot, day
                    break
            else:  # else clause for the 'for' statement
                extra_path.insert(0, NightlyPathEntry(base, slot, day))

    return extra_path


class LbRun(xenv.Script):
    __usage__ = ('Usage: %prog [OPTION]... [NAME=VALUE]... '
                 'PROJECT[/VERSION] [COMMAND [ARG]...]')

    def _prepare_parser(self):
        from options import addSearchPath, addPlatform, addListing
        from optparse import OptionValueError

        super(LbRun, self)._prepare_parser()
        parser = self.parser

        parser.add_option(
            "--use",
            action="append",
            help="add a data package to the environment")
        parser.add_option(
            "--ext",
            action="append",
            help="add an external lcg tool environment")

        addPlatform(parser)
        addSearchPath(parser)
        addListing(parser)

        def extract_project_version(opt_str, rargs):
            if not rargs:
                raise OptionValueError(
                    '%s must be followed by the project'
                    ' name and optionally by the version' % opt_str)
            if '/' in rargs[0]:
                p_name, v = rargs.pop(0).split('/')
                p_name = fixProjectCase(p_name)
            else:
                p_name = fixProjectCase(rargs.pop(0))
                if rargs and isValidVersion(p_name, rargs[0]):
                    v = rargs.pop(0)
                else:
                    v = DEFAULT_VERSION
            return p_name, v

        def runtime_project_option(_option, opt_str, _value, parser):
            pv = extract_project_version(opt_str, parser.rargs)
            parser.values.runtime_projects.append(pv)

        parser.add_option(
            "--runtime-project",
            action="callback",
            metavar="PROJECT[/VERSION]",
            type="string",
            callback=runtime_project_option,
            nargs=0,
            help="Add a project to the runtime environment")

        def overriding_project_option(_option, opt_str, _value, parser):
            pv = extract_project_version(opt_str, parser.rargs)
            parser.values.overriding_projects.append(pv)

        parser.add_option(
            "--overriding-project",
            action="callback",
            metavar="PROJECT[/VERSION]",
            type="string",
            callback=overriding_project_option,
            nargs=0,
            help="Add a project to override packages")

        parser.add_option(
            "--no-auto-override",
            action="store_false",
            dest="auto_override",
            help="Do not automatically prepend the "
            "projects %s" % auto_override_projects)

        parser.add_option(
            "--grid-override",
            action="store_true",
            dest='grid_override',
            help="Override grid middleware versions [default]")

        parser.add_option(
            "--no-grid-override",
            action="store_false",
            dest='grid_override',
            help="Prevent override grid middleware versions")

        parser.add_option(
            "--grid-override-map",
            help="Use specific grid middleware override map")

        parser.add_option(
            "--use-grid",
            action="store_true",
            help="ignored, kept for backward compatibility")

        parser.add_option(
            "--use-sp",
            action="store_true",
            dest="use_setupproject",
            help="Force fallback on SetupProject even if the "
            "project have manifest.xml and a .xenv files")

        parser.add_option(
            "--no-sp",
            action="store_true",
            dest="no_setupproject",
            help="Prevent fallback on SetupProject in case of "
            "problems")

        # Note: the profile is not used in the script class, but in the wrapper
        #       it is added to the parser to appear in the help and for
        #       checking
        parser.add_option(
            "--profile",
            action="store_true",
            help="Print some profile informations about the "
            "execution.")

        parser.add_option(
            '--path-to-project',
            action='store_true',
            help='Interpret the first argument as path to a '
            'project top-level directory instead of as a '
            'project name')

        parser.set_defaults(
            use=[],
            ext=[],
            runtime_projects=[],
            overriding_projects=[],
            auto_override=True,
            grid_override=True,
            use_grid=False,
            use_setupproject=False)
        self.allow_empty_version = False

    def _parse_args(self, args=None):
        super(LbRun, self)._parse_args(args)
        if len(self.cmd) < 1:
            self.parser.error("missing project " + (
                'path' if self.opts.path_to_project else 'name'))
        if self.opts.path_to_project:
            self.project = None
            self.version = os.path.abspath(self.cmd.pop(0))
        elif '/' in self.cmd[0]:
            self.project, self.version = self.cmd.pop(0).split('/')
            self.project = fixProjectCase(self.project)
        else:
            self.project = fixProjectCase(self.cmd.pop(0))
            if self.cmd and isValidVersion(self.project, self.cmd[0]):
                self.version = self.cmd.pop(0)
                self.log.warning('deprecated version specification: '
                                 'use "lb-run ... %s/%s ..." instead',
                                 self.project, self.version)
            else:
                # if no version is specified, we want to allow just the
                # project name
                self.allow_empty_version = True
                self.version = DEFAULT_VERSION
        if (self.opts.platform is None and
                str(self.project.lower()) == 'lhcbdirac'):
            self.opts.platform = 'best'
        else:
            from options import checkPlatform
            self.opts.platform = checkPlatform(self.parser, self.opts.platform)

    def _add_ext_dir_to_env(self, path):
        def prepend(k, v):
            self.log.debug('prepending %s to %s', v, k)
            self.opts.actions.append(('prepend', (k, v)))

        def with_python(files):
            for f in files:
                if f.endswith('.py'):
                    return True
            return False

        for root, dirs, files in os.walk(path):
            base = os.path.basename(root)
            if base == 'bin':
                prepend('PATH', os.path.join(root))
                dirs[:] = []  # no not recurse
            elif base.startswith('lib'):
                prepend('LD_LIBRARY_PATH', root)
                if with_python(files):
                    prepend('PYTHONPATH', root)
                # let recurse
            elif base.startswith('python'):
                if with_python(files):
                    prepend('PYTHONPATH', root)
                # let recurse
            elif base == 'site-packages':
                prepend('PYTHONPATH', root)
                dirs[:] = []  # no not recurse
            dirs[:] = [
                d for d in dirs
                if d in ['bin', 'site-packages', 'lib', 'lib64']
                or d.startswith('python')
            ]

    def _handle_externals(self, lcg_path, manifest):
        if not self.opts.ext:
            return
        if not lcg_path:
            self.log.warning('no LCG found: ignoring --ext option')
            return

        if manifest:
            _version, platform = getHepToolsInfo(manifest)
        else:
            platform = self.opts.platform
        ext_info_file = os.path.join(lcg_path, LCGInfoName(platform))
        if not os.path.isfile(ext_info_file):
            self.log.warning('no LCG_externals_<>.txt found: '
                             'ignoring --ext option')
            return

        def parse_deps_list(s):
            s = s.strip()
            if s:
                return [x.rsplit('-', 1)[0].lower() for x in s.split(',')]
            else:
                return []

        self.log.debug('handling required externals')
        # get list of ext dirs from LCG_externals_<platform>.txt
        # the result is a dictionary {name: (path, deps)}
        exts_info = dict(
            (x[0].lower(), (os.path.join(lcg_path, x[3]),
                            parse_deps_list(x[4])))
            for x in
            [map(str.strip, x.split(';')) for x in open(ext_info_file)]
            if len(x) == 5)

        # we want case insensitive lookup and the use may specify
        #  --ext A --ext B,C
        exts = set(
            map(str.lower, sum([x.split(',') for x in self.opts.ext], [])))
        if exts - set(exts_info):
            self.log.warning('ignored unknown externals: %s',
                             ', '.join(exts - set(exts_info)))
            exts = exts.intersection(exts_info)

        # expand dependencies
        count = -1
        while len(exts) != count:  # expand until we get the full list
            count = len(exts)  # size before expansion
            exts.update(sum([exts_info[ext][1] for ext in exts], []))

        if self.project == 'LCG':
            # for LCG we need to explicitly pick up the compiler
            for l in open(ext_info_file):
                if l.startswith('COMPILER:'):
                    name, version = l[9:].strip().split(';')
                    exts.add(name)
                    exts_info[name] = (os.path.join(
                        lcg_path, name, version, '-'.join(
                            platform.split('-')[:2])), None)
                    break
            else:  # this 'else' matches the 'for'
                raise NotFoundError('compiler in %s' % ext_info_file)

        # add all externals to the environment
        for ext in exts:
            path = exts_info[ext][0]
            if os.path.isdir(path):
                self._add_ext_dir_to_env(path)
                # special cases
                if ext == 'root':
                    self.opts.actions.append(('set', ('ROOTSYS', path)))
            else:
                self.log.warning('unusable path for %s: %s', ext, path)

    def _findBestPlatform(self):
        '''
        Return the "best" platform
         '''
        try:
            from itertools import ifilter
            platforms = listPlatforms(self.project, self.version)
            return ifilter(supported, platforms).next()

        except StopIteration:
            # a StopIteration from 'next()' means iterator empty
            self.log.error('current host does not support any of %s/%s '
                           'available platforms', self.project, self.version)
        raise SystemExit(64)

    def overrideGridMiddleware(self):
        '''
        Modify the environment to override middleware libraries.
        '''
        if not self.opts.grid_override:
            self.log.debug('not overriding grid middleware')
            return

        from json import load
        from re import sub, error as re_error
        from hashlib import sha1
        from itertools import chain

        try:
            if not self.opts.grid_override_map:
                self.opts.grid_override_map = os.path.join(
                    findDataPackage('LbEnvFix', 'prod'), 'override_map.json')
            override_map = load(open(self.opts.grid_override_map, 'rb'))

            substitutions = override_map.get('substitutions', [])

            checksum = sha1()
            for x in chain.from_iterable(substitutions):
                checksum.update(x)

            self.log.info('using override map version %s (sha1: %s)',
                          override_map.get('version', 'unknown'),
                          checksum.hexdigest())

            for varname in self.env:
                value = '{0}={1}'.format(varname, self.env[varname])
                for entry in substitutions:
                    a, b = entry[:2]
                    before = value
                    value = sub(a, b, value)
                    if value != before:
                        comment = '' if len(entry) < 3 else \
                                  ' ({})'.format(entry[2])
                        self.log.info('replaced %s with %s in %s%s', a, b,
                                      varname, comment)
                value = value.split('=', 1)
                if value[0] == varname:
                    self.env[varname] = value[1]
                else:
                    self.log.warning('variable rename not supported in '
                                     'override map (%s -> %s)', varname,
                                     value[0])
        except (NotFoundError, IOError, ValueError, re_error) as err:
            self.log.warning('%s: %s, not overriding grid middleware',
                             type(err).__name__, err)

    def _makeEnv(self):
        # FIXME: when we drop Python 2.4, this should become
        #        'from . import path'
        from LbEnv.ProjectEnv import path, SearchPathEntry

        if self.opts.use_grid:
            self.log.warning('the option --use-grid is ignored (deprecated)')

        # prepend dev dirs to the search path
        if self.opts.dev_dirs:
            path[:] = self.opts.dev_dirs + path

        if self.opts.user_area and not self.opts.no_user_area:
            path.insert(0, SearchPathEntry(self.opts.user_area))

        if self.opts.path_to_project:
            if self.opts.list:
                self.parser.error('options --list and --path-to-project '
                                  'are incompatible')
            elif self.opts.use_setupproject:
                self.parser.error('options --use-sp and --path-to-project '
                                  'are incompatible')

        # FIXME: we need to handle common options like --list in a single place
        if self.opts.list:
            for entry in listVersions(self.project, self.opts.platform):
                print '%s in %s' % entry
            sys.exit(0)
        if self.opts.list_platforms:
            platforms = listPlatforms(self.project, self.version)
            if platforms:
                print '\n'.join(platforms)
            sys.exit(0)

        if self.opts.platform.lower() == 'best':
            self.opts.platform = self._findBestPlatform()
        elif not supported(self.opts.platform):
            msg = ('current host does not support platform {0} '
                   '(dirac_platform: {1}, os_id: {2})').format(
                       self.opts.platform, HOST_PLATFORM,
                       LbPlatformUtils.os_id())
            if not self.opts.force_platform:
                self.log.error(msg)
                sys.exit(66)
            else:
                self.log.warning(msg)

        # special handling of external projects
        # (unless we only want to list)
        if self.project.upper() in EXTERNAL_PROJECTS:
            self.opts.ext.append(self.project)
            if self.allow_empty_version:  # no version specified
                # for LCG the version DEFAULT_VERSION does not make sense
                self.version = 'latest'
            else:
                self.version = findLCGForExt(self.project, self.version,
                                             self.opts.platform)
            self.project = 'LCG'

        self.version = expandVersionAlias(self.project, self.version,
                                          self.opts.platform)

        # prepare the list of projects to use
        projects = []
        if self.opts.auto_override:
            explicit = set([p[0] for p in self.opts.overriding_projects])
            projects.extend(
                [p for p in auto_override_projects if p[0] not in explicit])
        projects.extend(self.opts.overriding_projects)
        projects.append((self.project, self.version))
        projects.extend(self.opts.runtime_projects)

        # Check if the main project needs a special search path
        self.log.debug('check if we need extra search path')
        project_path = findProject(
            self.project,
            self.version,
            self.opts.platform,
            allow_empty_version=self.allow_empty_version)
        extra_path = projectExtraPath(project_path)
        if extra_path:
            self.log.debug('the project requires an extra search path')
            # we add the extra search path between the command line entries
            # and the default
            idx = len(self.opts.dev_dirs)
            if self.opts.user_area:
                idx += 1
                path[:] = path[:idx] + extra_path + path[idx:]
            self.log.debug('final search path: %r', path)

        # set the environment XML search path
        env_path = []
        for p, v in projects:
            if p == 'LCG':
                continue
            if p:
                v = expandVersionAlias(p, v, self.opts.platform)
            self.log.info('using %s/%s %s', p, v, self.opts.platform)
            env_path.extend(
                getEnvXmlPath(p, v, self.opts.platform,
                              self.allow_empty_version and p == self.project))

        # ensure that we do not have unicode strings
        # FIXME: xenv has got problems with unicode in the search path
        env_path = map(str, env_path)
        xenv.path.extend(env_path)

        # set LCG relocation roots
        if self.project == 'LCG':
            self.log.debug('project is LCG, using externals from %s',
                           project_path)
            lcg_relocation = {'LCG_external': project_path}
            manifest = None
        else:
            manifest = os.path.join(project_path, 'manifest.xml')
            lcg_relocation = getLCGRelocation(manifest)
        self.opts.actions.extend(
            ('set', (k, v)) for k, v in lcg_relocation.items())

        self._handle_externals(lcg_relocation.get('LCG_external'), manifest)

        # now we can expand project name and version if --path-to-project
        if self.opts.path_to_project:
            self.project, self.version = \
                getProjectNameVersion(os.path.join(project_path,
                                                   'manifest.xml'))
            # FIXME: xenv has got problems with unicode in the search path
            self.project, self.version = map(str, (self.project, self.version))

        # extend the prompt variable (bash, sh)
        if self.cmd and os.path.basename(self.cmd[0]) in ('bash', 'sh'):
            prompt = os.environ.get('PS1', r'\W \$ ')
        # extend the prompt variable (zsh)
        elif self.cmd and os.path.basename(self.cmd[0]) in ('zsh', ):
            prompt = os.environ.get('PS1', r'%1d%# ')
        else:
            prompt = '> '

        self.opts.actions.append(('set', ('PS1', r'[{0} {1}] {2}'.format(
            self.project, self.version, prompt))))

        # instruct the script to load the projects environment XML
        for p, _ in projects[::-1]:
            if p == 'LCG':
                continue
            if not p:  # this flags the main project (when --path-to-project)
                p = self.project
            self.opts.actions.append(('loadXML', (p + '.xenv', )))

        # handle the extra data packages
        for pkg_name, pkg_vers in map(decodePkg, self.opts.use):
            xml_name = pkg_name.replace('/', '_') + '.xenv'
            xml_path = os.path.join(
                findDataPackage(pkg_name, pkg_vers), xml_name)
            if not os.path.exists(xml_path):
                # fall back on the old conventional name
                xml_path = xml_path[:-5] + 'Environment.xml'
            # FIXME: xenv has got problems with unicode filenames
            self.opts.actions.append(('loadXML', (str(xml_path), )))

        # Set the CMTCONFIG and BINARY_TAG as they could different from the env
        self.opts.actions.append(('set', ('CMTCONFIG', self.opts.platform)))
        self.opts.actions.append(('set', ('BINARY_TAG', self.opts.platform)))

        try:
            super(LbRun, self)._makeEnv()
        except SystemExit as exc:
            # make sure that an exit from underlying _makeEnv has the right
            # bit set
            raise SystemExit(exc.code | 64)

        self.overrideGridMiddleware()

    def compatMain(self, reason=None):
        '''
        Fall-back function to call the old SetupProject code.
        '''
        from compatibility import getOldEnvironment
        if reason:
            self.log.warning('trying old SetupProject (%s)', reason)
        # fix command line to fit old SetupProject
        if self.cmd:
            args = sys.argv[1:-len(self.cmd)]
        else:
            args = sys.argv[1:]

        # if we ran "best platform match", pass the result to SetupProject
        if 'best' in args:
            args[args.index('best')] = self.opts.platform
        elif '-c=best' in args:
            args[args.index('-c=best')] = '-c=' + self.opts.platform
        elif '--platform=best' in args:
            args[args.index('--platform=best')] = '--platform=' + \
                self.opts.platform

        self.env = getOldEnvironment(args)

        self.overrideGridMiddleware()

        if self.cmd:
            return self.runCmd()
        else:
            self.dump()
            return 0

    def main(self):
        try:
            try:
                if not self.opts.use_setupproject:
                    super(LbRun, self).main()
                else:
                    sys.exit(self.compatMain())
            except (NotFoundError, IOError, OSError), x:
                if self.opts.path_to_project or self.opts.no_setupproject:
                    # SetupProject does not support --path-to-project
                    self.log.error('%s', x)
                    sys.exit(64)
                sys.exit(self.compatMain(x))
        except SystemExit:
            # pass through SystemExit exceptions (sys.exit)
            raise
        except:
            # force a special exit code for unhandled exceptions
            traceback.print_exc()
            raise SystemExit(65)


def main():
    # special handling of the option --profile
    if '--profile' in sys.argv:
        from LbEnv.ProjectEnv.profiling import run
        run('LbRun().main()')
    else:
        # special handling of the '#!' case (LBCORE-782)
        # - we must have at least one argument and the first one must be a file
        if len(sys.argv) >= 2 and os.path.isfile(sys.argv[1]):
            # - the second line of the file must start with '# args: '
            argline = ''
            with open(sys.argv[1]) as scriptfile:
                scriptfile.next()  # skip first line
                argline = scriptfile.next().strip()
            import re
            args = re.match(r'#\s*args\s*:\s*(\S+.*)', argline)
            if args:
                # FIXME: we should improve the argument splitting
                sys.argv = sys.argv[:1] + args.group(1).split() + sys.argv[1:]
        LbRun().main()
