#!/usr/bin/env python3

import os
import pathlib
import subprocess

import common
from shell_helpers import LF

class Main(common.BuildCliFunction):
    def __init__(self):
        super().__init__(
            description='''\
Build gem5.
https://github.com/cirosantilli/linux-kernel-module-cheat-regression#gem5-buildroot-setup
'''
        )
        self.add_argument(
            '--unit-test',
            action='append',
            default=[],
            help='''\
Build and run the given unit test. Paths are relative to src/ without the .opt suffix.
If given multiple times, runs multiple unit tests. Ignore --unit-tests.
https://github.com/cirosantilli/linux-kernel-module-cheat-regression#gem5-unit-tests
'''
        )
        self.add_argument(
            '--unit-tests',
            default=False,
            help='''\
Build and run all the gem5 unit tests instead of the gem5 executable.
https://github.com/cirosantilli/linux-kernel-module-cheat-regression#gem5-unit-tests
'''
        )
        self.add_argument(
            'extra_scons_args',
            metavar='extra-scons-args',
            nargs='*',
        )

    def build(self):
        build_dir = self.get_build_dir()
        binaries_dir = os.path.join(self.env['gem5_system_dir'], 'binaries')
        disks_dir = os.path.join(self.env['gem5_system_dir'], 'disks')
        os.makedirs(binaries_dir, exist_ok=True)
        os.makedirs(disks_dir, exist_ok=True)
        if not os.path.exists(os.path.join(self.env['gem5_source_dir'], '.git')):
            if self.env['_args_given']['gem5_worktree']:
                self.sh.run_cmd([
                    'git', LF,
                    '-C', self.env['gem5_default_source_dir'], LF,
                    'worktree', 'add', LF,
                    '-b', os.path.join('wt', self.env['gem5_worktree']), LF,
                    self.env['gem5_source_dir'], LF,
                ])
            else:
                if not self.env['dry_run']:
                    raise Exception('gem5 submodule not checked out')
        if self.env['verbose']:
            verbose = ['--verbose', LF]
        else:
            verbose = []
        if self.env['arch'] == 'x86_64':
            dummy_img_path = os.path.join(disks_dir, 'linux-bigswap2.img')
            with open(dummy_img_path, 'wb') as dummy_img_file:
                zeroes = b'\x00' * (2 ** 16)
                for i in range(2 ** 10):
                    dummy_img_file.write(zeroes)
            self.sh.run_cmd(['mkswap', dummy_img_path, LF])
            with open(os.path.join(binaries_dir, 'x86_64-vmlinux-2.6.22.9'), 'w'):
                # This file must always be present, despite --kernel overriding that default and selecting the kernel.
                # I'm not even joking. No one has ever built x86 gem5 without the magic dist dir present.
                pass
        elif self.env['arch'] == 'arm' or self.env['arch'] == 'aarch64':
            gem5_system_source_dir = os.path.join(self.env['gem5_source_dir'], 'system')

            # dtb
            dt_source_dir = os.path.join(gem5_system_source_dir, 'arm', 'dt')
            dt_build_dir = os.path.join(self.env['gem5_system_dir'], 'arm', 'dt')
            self.sh.run_cmd(['make', '-C', dt_source_dir, LF])
            self.sh.copy_dir_if_update_non_recursive(
                srcdir=dt_source_dir,
                destdir=dt_build_dir,
                filter_ext='.dtb',
            )

            # Bootloader 32.
            bootloader32_dir = os.path.join(gem5_system_source_dir, 'arm', 'simple_bootloader')
            # TODO use the buildroot cross compiler here, and remove the dependencies from configure.
            self.sh.run_cmd([
                'make', LF,
                '-C', bootloader32_dir, LF,
                'CROSS_COMPILE=arm-linux-gnueabihf-', LF,
            ])
            # bootloader
            self.sh.cp(os.path.join(bootloader32_dir, 'boot_emm.arm'), binaries_dir)

            # Bootloader 64.
            bootloader64_dir = os.path.join(gem5_system_source_dir, 'arm', 'aarch64_bootloader')
            # TODO cross_compile is ignored because the make does not use CC...
            self.sh.run_cmd(['make', '-C', bootloader64_dir, LF])
            self.sh.cp(os.path.join(bootloader64_dir, 'boot_emm.arm64'), binaries_dir)
            self.sh.cp(os.path.join(bootloader64_dir, 'boot_emm_v2.arm64'), binaries_dir)
        term_source_dir = os.path.join(self.env['gem5_source_dir'], 'util/term')
        m5term_build = os.path.join(term_source_dir, 'm5term')
        self.sh.run_cmd(['make', '-C', term_source_dir, LF])
        if os.path.exists(self.env['gem5_m5term']):
            # Otherwise self.sh.cp would fail with "Text file busy" if you
            # tried to rebuild while running m5term:
            # https://stackoverflow.com/questions/16764946/what-generates-the-text-file-busy-message-in-unix/52427512#52427512
            self.sh.rmrf(self.env['gem5_m5term'])
        self.sh.cp(m5term_build, self.env['gem5_m5term'])
        if self.env['unit_test']:
            targets = [self.get_gem5_target_path(self.env, test) for test in self.env['unit_test']]
        elif self.env['unit_tests']:
            targets = [self.env['gem5_unit_test_target']]
        else:
            targets = [self.env['gem5_executable']]
        if self.env['gem5_clang']:
            extra_env = {
                'CC': 'clang',
                'CXX': 'clang++',
            }
            gold_linker_cmd = []
        else:
            extra_env = {}
            gold_linker_cmd = ['--gold-linker', LF,]
        exit_status = self.sh.run_cmd(
            (
                [
                    'scons', LF,
                    '-j', str(self.env['nproc']), LF,
                    '--ignore-style', LF,
                ] +
                gold_linker_cmd +
                verbose +
                [
                    'SLICC_HTML=True', LF,
                ] +
                self.sh.add_newlines(targets) +
                self.sh.add_newlines(self.env['extra_scons_args'])
            ),
            cwd=self.env['gem5_source_dir'],
            extra_paths=[self.env['ccache_dir']],
            extra_env=extra_env,
            raise_on_failure = False,
        )
        return exit_status

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

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