#!/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
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 Mu(Language):
    name = 'mu'
    src_ext = 'c'  # one or a folder of C files calling API to build the bundle
    compiled = True
    default_exec = 'clang'  # assume presence of clang
    known_impls = ('holstein', 'zebu')

    @classmethod
    def check_lang(cls, lc):
        assert 'impl' in lc, "Mu impl not specified"
        assert lc['impl'] in cls.known_impls, 'invalid impl: %(impl)s' % lc
        return lc

    @classmethod
    def check_compiler(cls, cc, lc, task):
        impl = lc['impl']
        env = task.env

        # check Mu impl environment variable
        v = 'MU_%s' % (impl.upper())
        assert v in env, '%s needs to be defined in environ.'

        mudir = Path(env[v])

        exe = cc.get('exec', cls.default_exec)
        cc['exec'] = expandenv(exe, env)

        # include dirs
        include_dirs = cc.get('include_dirs', [])
        include_dirs.append(CALLBACKS_DIR / cls.name)
        if impl == 'holstein':
            include_dirs.append(mudir / 'cbinding')
        elif impl == 'zebu':
            include_dirs.append(mudir / 'src' / 'vm' / 'api')
        cc['include_dirs'] = include_dirs

        # library dirs
        library_dirs = cc.get('library_dirs', [])
        libmu_dir = ''
        libname = ''
        if impl == 'holstein':
            libmu_dir = mudir / 'cbinding'
            libname = 'murefimpl2start'
        elif impl == 'zebu':
            libmu_dir = mudir / 'target' / env.get('ZEBU_BUILD', 'release')
            libname = 'mu'
        library_dirs.append(libmu_dir)
        cc['library_dirs'] = library_dirs
        cc.setdefault('libmu_link_name', libname)

        # vmarg
        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 key in vmargs:
                    vmargs.append('%(key)s=%(val)s' % locals())
            return vmargs

        vmargs = list(map(lambda a: expandenv(a, task.env), cc.get('vmargs', [])))
        libcb = task.taskset.callback['dylib']
        if lc['impl'] == 'holstein':
            separators = '\n'
            default_args = {
                'extraLibs': str(libcb)
            }
        else:   # zebu
            separators = ' '
            emit_dir = task.output_dir / ('%s_%s-emit' % (task.taskset.name, task.name))
            default_args = {
                '--aot-emit-dir': str(emit_dir),
                '--bootimage-external-libpath': str(libcb.parent),
                '--bootimage-external-lib': "cb_%(name)s" % task.callback   # libcb_%(name)s.so/dylib
            }
        vmarg_s = separators.join(set_default_args(default_args, vmargs))
        if lc['impl'] == 'zebu':
            vmarg_s = 'init_mu ' + vmarg_s
        cc['vmarg_s'] = vmarg_s
        return cc

    @classmethod
    def check_runner(cls, rc, lc, task):
        env = task.env
        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',
                '--extraLibs=%s' % task.taskset.callback['dylib']]
            rc['flags'] = holstein_default_flags + rc.get('flags', [])
        return rc

    @classmethod
    def compile(cls, task):
        callback_dir = CALLBACKS_DIR / cls.name

        cc = task.compiler
        lc = task.lang

        # First compile the build scripts
        cmd = []
        cmd.append(cc['exec'])

        # flags
        flags = []

        impl = lc['impl']
        flags.append('-DMU_IMPL_%s' % impl.upper())

        flags.extend(['-I%s' % d for d in cc['include_dirs']])
        flags.extend(['-L%s' % d for d in cc['library_dirs']])
        flags.append('-l%(libmu_link_name)s' % cc)

        target = cls.get_default_target(task)
        build_target = target.with_name('build_' + target.name)
        flags.extend(['-o', str(build_target)])

        cmd.extend(flags)

        bm_src = task.srcfile
        srcs = [
            callback_dir / 'build_callbacks.c',
            callback_dir / 'main.c',
            bm_src
        ]
        cmd.extend(srcs)

        cls.run_in_subproc(cmd, task.env)

        assert build_target.exists()

        # Then run the build program to get the boot image

        cmd = [build_target]
        cmd.append(cc['vmarg_s'])
        cmd.append(target.absolute())

        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['impl'] == 'holstein':
            cmd.append(rc['exec'])
            cmd.extend(rc.get('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)