#!/usr/bin/env python3

import os
import shlex

import common
import threading
import subprocess
from shell_helpers import LF

error = False

class Main(common.BuildCliFunction):
    def __init__(self):
        super().__init__(
            description='''\
Build our compiled userland examples.
'''
        )
        self.add_argument(
            '--has-package',
            action='append',
            default=[],
            help='''\
Indicate that a given package is present in the root filesystem, which
allows us to build examples that rely on it.
''',
        )
        self.add_argument(
            'targets',
            default=[],
            help='''\
Build only the given userland programs.
Default: build all examples that have their package dependencies met.
For example, an OpenBLAS example can only be built if the target root filesystem
has the OpenBLAS libraries and headers installed.
''',
            nargs='*',
        )

    def _build_one(
        self,
        in_path,
        out_path,
        ccflags,
        extra_deps=None,
        link=True,
        extra_objs=None,
        std=None,
        ccflags_after=None,
        raise_on_failure=True,
        thread_limiter=None,
    ):
        try:
            if extra_deps is None:
                extra_deps = []
            if extra_objs is None:
                extra_objs = []
            if ccflags_after is None:
                ccflags_after = []
            ret = 0
            if self.need_rebuild([in_path] + extra_objs + extra_deps, out_path):
                ccflags = ccflags.copy()
                if not link:
                    ccflags.extend(['-c', LF])
                in_ext = os.path.splitext(in_path)[1]
                do_compile = True
                if in_ext == self.env['c_ext']:
                    cc = self.env['gcc']
                    if std is None:
                        std = 'c11'
                    ccflags.extend([
                        '-fopenmp', LF,
                    ])
                elif in_ext == self.env['cxx_ext']:
                    cc = self.env['gxx']
                    if std is None:
                        std = 'c++17'
                else:
                    do_compile = False
                if do_compile:
                    ret = self.sh.run_cmd(
                        (
                            [
                                cc, LF,
                            ] +
                            ccflags +
                            [
                                '-std={}'.format(std), LF,
                                '-o', out_path, LF,
                                in_path, LF,
                            ] +
                            extra_objs +
                            [
                                '-lm', LF,
                                '-pthread', LF,
                            ] +
                            ccflags_after
                        ),
                        extra_paths=[self.env['ccache_dir']],
                        raise_on_failure=raise_on_failure,
                    )
        finally:
            if thread_limiter is not None:
                thread_limiter.release()
        if ret != 0:
            self.error = True
        return ret

    def build(self):
        build_dir = self.get_build_dir()
        os.makedirs(build_dir, exist_ok=True)
        has_packages = set(self.env['has_package'])
        ccflags = [
            '-I', self.env['root_dir'], LF,
            '-I', self.env['userland_source_dir'], LF,
            '-O0', LF,
            '-Wall', LF,
            '-Werror', LF,
            '-Wextra', LF,
            '-Wno-unused-function', LF,
            '-pedantic', LF,
            '-ggdb3', LF,
        ]
        if self.env['static']:
            ccflags.extend(['-static', LF])
        common_obj = os.path.join(
            build_dir,
            self.env['common_basename_noext'] + self.env['obj_ext']
        )
        self._build_one(
            in_path=self.env['common_c'],
            out_path=common_obj,
            ccflags=ccflags,
            extra_deps=[self.env['common_h']],
            link=False,
        )
        pkgs = {
            'eigen': {
                # TODO: was failing with:
                # fatal error: Eigen/Dense: No such file or directory as of
                # 975ce0723ee3fa1fea1766e6683e2f3acb8558d6
                # http://lists.busybox.net/pipermail/buildroot/2018-June/222914.html
                'ccflags': [
                    '-I',
                    os.path.join(
                        self.env['buildroot_staging_dir'],
                        'usr',
                        'include',
                        'eigen3'
                    ),
                    LF
                ],
                # Header only.
                'ccflags_after': [],
            },
            'libdrm': {},
            'openblas': {},
        }
        rootdir_abs_len = len(self.env['userland_source_dir'])
        thread_limiter = threading.BoundedSemaphore(self.env['nproc'])
        self.error = False
        for path, in_dirnames, in_filenames in os.walk(self.env['userland_source_dir']):
            in_dirnames.sort()
            dirpath_relative_root = path[rootdir_abs_len + 1:]
            dirpath_relative_root_components = dirpath_relative_root.split(os.sep)
            if (
                len(dirpath_relative_root_components) < 2 or
                dirpath_relative_root_components[0] != 'arch' or
                dirpath_relative_root_components[1] == self.env['arch']
            ):
                out_dir = os.path.join(
                    self.env['userland_build_dir'],
                    dirpath_relative_root
                )
                os.makedirs(out_dir, exist_ok=True)
                for in_filename in in_filenames:
                    in_path = os.path.join(path, in_filename)
                    in_name, in_ext = os.path.splitext(in_filename)
                    out_path = os.path.join(
                        out_dir,
                        in_name + self.env['userland_build_ext']
                    )
                    pkg_key = in_name.split('_')[0]
                    ccflags_file = ccflags.copy()
                    ccflags_after = []
                    if pkg_key in pkgs:
                        if pkg_key not in has_packages:
                            continue
                        pkg = pkgs[pkg_key]
                        if 'ccflags' in pkg:
                            ccflags_file.extend(pkg['ccflags'])
                        else:
                            pkg_config_output = subprocess.check_output([
                                self.env['buildroot_pkg_config'],
                                '--cflags',
                                pkg_key
                            ]).decode()
                            ccflags_file.extend(self.sh.shlex_split(pkg_config_output))
                        if 'ccflags_after' in pkg:
                            ccflags_file.extend(pkg['ccflags_after'])
                        else:
                            pkg_config_output = subprocess.check_output([
                                self.env['buildroot_pkg_config'],
                                '--libs',
                                pkg_key
                            ]).decode()
                            ccflags_after.extend(self.sh.shlex_split(pkg_config_output))
                    thread_limiter.acquire()
                    if self.error:
                        return 1
                    thread = threading.Thread(
                        target=self._build_one,
                        kwargs={
                            'in_path':in_path,
                            'out_path':out_path,
                            'ccflags':ccflags_file,
                            'extra_objs':[common_obj],
                            'ccflags_after':ccflags_after,
                            'raise_on_failure':False,
                            'thread_limiter':thread_limiter,
                        }
                    )
                    thread.start()
        self.sh.copy_dir_if_update(
            srcdir=build_dir,
            destdir=self.env['out_rootfs_overlay_dir'],
            filter_ext=self.env['userland_build_ext'],
        )
        return 0

    def get_build_dir(self):
        return self.env['userland_build_dir']

if __name__ == '__main__':
    Main().cli()