#!/usr/bin/env python3

import argparse
from pathlib import Path
import multiprocessing
import itertools
from concurrent.futures import ThreadPoolExecutor
from typing import Sequence, Iterable, Dict, Tuple
import subprocess
import time
import sys

HERE_DIR = Path(__file__).parent.absolute()

INCLUDE_DIRS = [
    'external/taywee-args/include',
    'external/spdlog/include',
]


def _compile_src(cxx: Path, cpp_file: Path) -> Tuple[Path, Path]:
    build_dir = HERE_DIR / '_build'
    src_dir = HERE_DIR / 'src'
    relpath = cpp_file.relative_to(src_dir)
    obj_path = build_dir / relpath.with_name(relpath.name + '.o')
    obj_path.parent.mkdir(exist_ok=True, parents=True)
    cmd = [
        cxx,
        '-std=c++17',
        '-static',
        '-Wall',
        '-Wextra',
        '-Werror',
        '-Wshadow',
        '-Wconversion',
        '-fdiagnostics-color',
        '-pthread',
        '-g',
        '-c',
        '-O0',
        f'-I{src_dir}',
        cpp_file,
        f'-o{obj_path}',
    ]
    cmd.extend(
        itertools.chain.from_iterable(
            ('-isystem', HERE_DIR / subdir) for subdir in INCLUDE_DIRS))
    msg = f'Compile C++ file: {cpp_file}'
    print(msg)
    start = time.time()
    res = subprocess.run(
        cmd,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
    )
    if res.returncode != 0:
        raise RuntimeError(
            f'Compile command ({cmd}) failed for {cpp_file}:\n{res.stdout.decode()}'
        )
    end = time.time()
    print(f'{msg} - Done: {end - start:.2}s')
    return cpp_file, obj_path


def compile_sources(cxx: Path, sources: Iterable[Path]) -> Dict[Path, Path]:
    pool = ThreadPoolExecutor(multiprocessing.cpu_count() + 2)
    return {
        src: obj
        for src, obj in pool.map(lambda s: _compile_src(cxx, s), sources)
    }


def make_library(objects: Iterable[Path]) -> Path:
    lib_file = HERE_DIR / '_build/libddslim.a'
    cmd = ['ar', 'rsc', lib_file]
    cmd.extend(objects)
    if lib_file.exists():
        lib_file.unlink()
    print(f'Creating static library {lib_file}')
    subprocess.check_call(cmd)
    return lib_file


def link_exe(cxx: Path, obj: Path, lib: Path, *, out: Path = None) -> Path:
    if out is None:
        basename = obj.stem
        out = HERE_DIR / '_build/test' / (basename + '.exe')
        out.parent.mkdir(exist_ok=True, parents=True)

    print(f'Linking executable {out}')
    subprocess.check_call([
        cxx,
        '-static',
        '-pthread',
        # See: https://stackoverflow.com/questions/35116327/when-g-static-link-pthread-cause-segmentation-fault-why
        '-Wl,--whole-archive',
        '-lpthread',
        '-Wl,--no-whole-archive',
        obj,
        lib,
        '-lstdc++fs',
        f'-o{out}',
    ])
    return out


def run_test(exe: Path) -> None:
    print(f'Running test: {exe}')
    subprocess.check_call([exe])


def main(argv: Sequence[str]) -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--test', action='store_true', help='Build and run tests')
    parser.add_argument(
        '--cxx', help='Path/name of the C++ compiler to use.', required=True)
    args = parser.parse_args(argv)

    all_sources = set(HERE_DIR.glob('src/**/*.cpp'))
    test_sources = set(HERE_DIR.glob('src/**/*.test.cpp'))
    main_sources = set(HERE_DIR.glob('src/**/*.main.cpp'))

    lib_sources = (all_sources - test_sources) - main_sources

    objects = compile_sources(Path(args.cxx), all_sources)

    lib = make_library(objects[p] for p in lib_sources)

    test_objs = (objects[p] for p in test_sources)
    pool = ThreadPoolExecutor(multiprocessing.cpu_count() + 2)
    test_exes = list(
        pool.map(lambda o: link_exe(Path(args.cxx), o, lib), test_objs))

    main_exe = link_exe(
        Path(args.cxx),
        objects[next(iter(main_sources))],
        lib,
        out=HERE_DIR / '_build/ddslim')

    if args.test:
        list(pool.map(run_test, test_exes))

    print(f'Main executable generated at {main_exe}')
    return 0


if __name__ == "__main__":
    sys.exit(main(sys.argv[1:]))