#!/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.evn

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

        mu_dir = 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(mu_dir / 'cbinding')
        elif impl == 'zebu':
            include_dirs.append(mu_dir / 'src' / 'vm' / 'api')
        cc['include_dirs'] = include_dirs

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

        return cc

    @classmethod
    def check_runner(cls, rc, lc, 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']
            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_dir = task.output_dir
        build_target = target_dir / task.resfile.name[:-2]
        flags.extend(['-o', str(build_target)])

        cmd.extend(flags)

        bm_src = task.resfile
        srcs = [
            callback_dir / 'build_callbacks.c',
            callback_dir / 'main.c',
            callback_dir / ('build_cb_%(name)s.c' % task.callback),
            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
        if bm_src.name.startswith('build_'):
            target_name = '%s-%s-mu' % (bm_src.name[len('build_'):-2],
                                        task.callback['name'])
        else:
            target_name = 'bootimg-' + bm_src.name[:-2]
        target = target_dir / target_name
        cmd = [build_target]
        cmd.append(cc.get('vmarg', ""))
        cmd.append(target)

        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 'exec' in rc:
            cmd.append(rc['exec'])
            cmd.extend(rc.get('flags', []))

        if lc['impl'] == 'holstein':
            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)