#!/usr/bin/env python3
# Copyright 2017 The Australian National University
#
# 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.

import subprocess as subproc
import os
import sys
from pathlib import Path

from mubench.lang import Language
from mubench import CALLBACKS_DIR
from mubench.exceptions import ExecutionFailure
from mubench.util import expandenv


class RPython(Language):
    name = 'rpython'
    src_ext = 'py'
    compiled = True
    default_python = 'pypy'  # use pypy by default, it's faster
    default_flags = ['-O3', '--no-shared']
    known_impls = ('holstein', 'zebu')

    @classmethod
    def check_lang(cls, lc):
        assert 'backend' in lc, "backend not specified."
        backend = lc['backend']
        assert backend in ('c', 'mu'), \
            "invalid backend: %(backend)s." % locals()
        if backend == 'mu':
            assert 'impl' in lc, "Mu impl needs to be specified for Mu backend."
            assert lc['impl'] in cls.known_impls, 'invalid impl: %(impl)s' % lc
        return lc

    @classmethod
    def check_compiler(cls, cc, lc, task):
        env = task.env

        # python interpreter
        python = cc.get('exec', cls.default_python)
        cc['exec'] = expandenv(python, env)

        # rpython
        assert 'PYPY' in env, 'PYPY needs to be defined in environ'
        pypydir = Path(env['PYPY'])
        rpython_py = pypydir / 'rpython' / 'bin' / 'rpython'
        assert rpython_py.exists(), '%(rpython_py)s not found' % locals()
        cc['rpython'] = rpython_py

        # set PYTHONPATH
        # include: pypydir, callbacks/rpython
        pythonpath = env.get('PYTHONPATH', '')
        callback_dir = CALLBACKS_DIR / cls.name
        pythonpath = '%(pypydir)s:%(callback_dir)s:%(pythonpath)s' % locals()
        env['PYTHONPATH'] = pythonpath

        # check Mu impl environment variable
        if lc['backend'] == 'mu':
            v = 'MU_%s' % (lc['impl'].upper())
            assert v in env, '%s needs to be defined in environ.'

        # target configuration arguments
        cc['args'] = list(map(lambda a: expandenv(str(a), task.env), cc.get('args', [])))

        # flags
        cc['flags'] = list(map(lambda a: expandenv(a, task.env), cc.get('flags', [])))

        if lc['backend'] == 'mu':
            # vmargs
            def set_default_args(defl_d, vmargs):
                # include default args if not specified, otherwise use custom-defined value
                for key, val in defl_d.items():
                    if not any([key in arg for arg in vmargs]):
                        vmargs.append('%(key)s=%(val)s' % locals())
                return vmargs

            vmargs = list(map(lambda a: expandenv(a, task.env), cc.get('vmargs', [])))
            if lc['impl'] == 'holstein':
                separators = '\n'
                default_args = {}
            else:   # zebu
                separators = ' '
                emit_dir = task.output_dir / ('%s_%s-emit' % (task.taskset.name, task.name))
                default_args = {
                    '--aot-emit-dir': str(emit_dir)
                }
            vmargs = set_default_args(default_args, vmargs)
            vmarg_s = separators.join(vmargs)
            cc['vmarg_s'] = vmarg_s

        return cc

    @classmethod
    def check_runner(cls, rc, lc, task):
        env = task.env
        if lc['backend'] == 'mu':
            suplibdir = str(task.output_dir)
            if sys.platform == 'darwin':
                env['DYLD_LIBRARY_PATH'] = suplibdir
            else:
                env['LD_LIBRARY_PATH'] = suplibdir

            if lc['impl'] == 'holstein':
                mu_holstein_dir = Path(env['MU_HOLSTEIN'])
                default_exec = mu_holstein_dir / 'tools' / 'runmu.sh'
                rc.setdefault('exec', default_exec)
                rc['exec'] = Path(expandenv(str(rc['exec']), env))

                holstein_default_flags = ['--uPtrHack=True']
                rc_flags = list(map(lambda a: expandenv(a, task.env), rc.get('flags', [])))
                rc['flags'] = holstein_default_flags + rc_flags

        return rc

    @classmethod
    def compile(cls, task):
        cmd = []

        cmd.append(task.compiler['exec'])
        cmd.append(task.compiler['rpython'])

        # flags
        bk = task.lang['backend']
        flags = []
        flags.extend(cls.default_flags)
        flags.append('--backend=%(bk)s' % locals())
        if bk == 'mu':
            flags.append('--mu-impl=%(impl)s' % task.lang)
            flags.append('--mu-suplibdir=%s' % task.output_dir)
            flags.append("--mu-vmargs=%(vmarg_s)s" % task.compiler)

        flags.extend(task.compiler['flags'])
        target = cls.get_default_target(task)
        flags.append('--output=%(target)s' % locals())
        cmd.extend(flags)

        cmd.append(task.srcfile)

        # target configuration arguments
        # first argument: callback lib
        cmd.append(task.callback['dylib'])

        # the rest, custom defined arguments
        cmd.extend(task.compiler['args'])

        cls.run_in_subproc(cmd, task.env)

        assert target.exists()
        return target

    @classmethod
    def run(cls, target, task):
        cmd = []
        rc = task.runner
        lc = task.lang

        if lc['backend'] == 'mu' and lc['impl'] == 'holstein':
            cmd.append(rc['exec'])
            cmd.extend(rc['flags'])

        cmd.append(target)

        # first argument: callback param
        cmd.append(task.callback['param'])

        cmd.extend(task.benchmark['args'])

        return cls.run_in_subproc(cmd, task.env)