#! python3

from argparse import ArgumentParser
import hashlib
import json
import os
from pathlib import Path
import subprocess
import shutil
import sys


class Config:
    CC = 'g++'                             # Compiler to call
    APPS = ['app_name']                    # Output binaries (need to be found as .cpp directly in SOURCE_DIR)
    IGNORE_APPS = []
    JOB_COUNT = int(os.cpu_count() * 0.8)  # Concurent jobs (multi-processing)
    WATCH = False                          # Watch source modification for auto-compiling

    BIN_DIR = Path('bin')                # Output directory (binaries)
    INCLUDE_DIR = Path('include')  # Include directory (header files)
    OBJECT_DIR = Path('obj')       # Temporary directory (object files)
    SOURCE_DIR = Path('src')       # Source directories

    COMMON_FLAGS = '-std=c++17'               # Flags used for comiling and linking
    COMMON_DEBUG_FLAGS = '-g'                 # Flags added in debug mode
    COMMON_RELEASE_FLAGS = '-O2 -flto'        # Flags added in release mode
    COMPILE_FLAGS = f'-Wall -I{INCLUDE_DIR}'  # Flags added for compiling (recommandation : `pkg-config --cflags`)
    LINK_FLAGS = ''                           # Flags added for linking (recommandation : `pkg-config --libs`)

    PRE_COMPILE_FUNCTION = None  # Function to run before compile (not call in --clean situation)

    CPP_SOURCES = SOURCE_DIR.rglob('*.cpp')


class ConsoleColor:
    """Simple shortcut to use colors in console."""

    HEADER = '\033[95m'
    BLUE = '\033[94m'
    GREEN = '\033[92m'
    ORANGE = '\033[93m'
    RED = '\033[91m'
    ENDCOLOR = '\033[0m'
    BOLD = '\033[1m'
    UNDERLINE = '\033[4m'


def get_hash(path: Path) -> str:
    hash_obj = hashlib.md5()
    with open(path, 'r', encoding='utf-8') as hashing_file:
        hash_obj.update(hashing_file.read().encode())
    return hash_obj.hexdigest()


def make(config: Config):
    parser = ArgumentParser()
    parser.add_argument('-j', type=int, default=config.JOB_COUNT, help='Jobs count (multi-processing)')
    parser.add_argument('--type', default='production', help='Compilation type (release, debug). Default=release')
    parser.add_argument('--clean', action='store_true', help='Clean all file instead of building')
    arguments = parser.parse_args()

    # Clean action
    if arguments.clean:
        if config.OBJECT_DIR.exists():
            for object_entry in config.OBJECT_DIR.iterdir():
                shutil.rmtree(object_entry, ignore_errors=True)
        if config.BIN_DIR.exists():
            for binary_entry in config.BIN_DIR.iterdir():
                shutil.rmtree(binary_entry, ignore_errors=True)
        return

    if config.PRE_COMPILE_FUNCTION is not None:
        config.PRE_COMPILE_FUNCTION()
    if 'IGNORE_APPS' in config.__dict__:
        config.IGNORE_APPS = []

    # Update flags and directories for mode debug/release
    if arguments.type == 'debug':
        config.COMMON_FLAGS += ' ' + config.COMMON_DEBUG_FLAGS
        config.OBJECT_DIR = config.OBJECT_DIR / 'debug'
    else:
        config.COMMON_FLAGS += ' ' + config.COMMON_RELEASE_FLAGS
        config.OBJECT_DIR = config.OBJECT_DIR / 'release'

    # Update job count
    config.JOB_COUNT = arguments.j

    Builder(config).make()


class Builder:
    def __init__(self, config: Config):
        self._config = config

        # Create list of source to process (tuple[source_path, object_path])
        self._compile_dict: dict[Path, Path] = {
            Path(source_file): (self._config.OBJECT_DIR / source_file.parent.relative_to(self._config.SOURCE_DIR)
                                / (source_file.stem + '.o'))
            for source_file in self._config.CPP_SOURCES}
        self._todo_dict: dict[Path, str] = {}

        self._hash_dict: dict[str, str] = {}
        if (self._config.OBJECT_DIR / 'hash.json').exists():
            with open(self._config.OBJECT_DIR / 'hash.json', 'r', encoding='utf-8') as hash_file:
                self._hash_dict = json.loads(hash_file.read())

        # Get source dependencies
        error_paths: list[tuple[Path, str]] = []
        self._dependency_dict: dict[str, list[str]] = {}
        for source_path, _object_path in self._compile_dict.items():
            cmd = ' '.join(
                [self._config.CC, self._config.COMMON_FLAGS, self._config.COMPILE_FLAGS, str(source_path), '-M'])
            job = subprocess.run(cmd, check=False, stdout=subprocess.PIPE, stderr=subprocess.PIPE,
                                 universal_newlines=True, shell=True)
            if job.returncode != 0:
                error_paths.append((source_path, job.stdout + job.stderr))
                break
            self._dependency_dict[str(source_path)] = [
                line for line in job.stdout.split('.o: ')[1].replace('\n', '').replace('\\ ', '').split(' ') if line]
        if error_paths:
            for error_path, error_text in error_paths:
                print(f'{ConsoleColor.RED}Error checking dependencies for {error_path}:'
                      f'\n{ConsoleColor.ENDCOLOR + error_text}')
            with open(self._config.OBJECT_DIR / 'hash.json', 'w', encoding='utf-8') as hash_file:
                hash_file.write(json.dumps(self._hash_dict, indent=1))
            sys.exit(1)

    def make(self):
        self._populate_todo_dict()

        if self._config.WATCH:
            from watch import Watcher, WatchFlag  # pylint: disable=import-outside-toplevel

            def callback(path: Path, _flags: WatchFlag):
                source_hash = get_hash(path)
                if str(path) not in self._hash_dict or source_hash != self._hash_dict[str(path)]:
                    print(f'{ConsoleColor.ORANGE}Source file changed : {path} {ConsoleColor.ENDCOLOR}')
                    self._populate_todo_dict()
                    self._compile_sources()

            self._compile_sources()
            watcher = Watcher()
            for source_path in self._compile_dict:
                watcher.register(source_path, WatchFlag.MODIFY, callback)
            try:
                watcher.watch()
            except KeyboardInterrupt:
                print('\rExit')
        elif not self._compile_sources():
            sys.exit(1)

    def _populate_todo_dict(self):
        """Check source file and generate compilation commands to execute."""
        for source_path, object_path in self._compile_dict.items():
            if not object_path.parent.exists():
                object_path.parent.mkdir(parents=True)
            source_hash = get_hash(source_path)
            dependency_changed = False
            for path in self._dependency_dict[str(source_path)]:
                dependency_hash = get_hash(path)
                if str(path) not in self._hash_dict or self._hash_dict[str(path)] != dependency_hash:
                    dependency_changed = True
                    print(f'{ConsoleColor.ORANGE}Dependency changed for {source_path} : {path} {ConsoleColor.ENDCOLOR}')
                    break
            if (dependency_changed or not object_path.exists()
                    or str(source_path) not in self._hash_dict or self._hash_dict[str(source_path)] != source_hash):
                self._todo_dict[source_path] = self._generate_compilation_command(source_path, object_path)
                continue

    def _generate_compilation_command(self, source_path: Path, object_path: Path) -> str:
        return ' '.join([self._config.CC, self._config.COMMON_FLAGS, self._config.COMPILE_FLAGS,
                        str(source_path), '-c', '-o', str(object_path)])

    def _compile_sources(self) -> bool:
        if not self._todo_dict and all(
                (self._config.BIN_DIR / app_path).exists()
                for app_path in self._config.APPS if app_path not in self._config.IGNORE_APPS):
            print(f'{ConsoleColor.GREEN}Nothing to do{ConsoleColor.ENDCOLOR}')
            return True

        # Running compilation processes
        if not os.path.exists(self._config.OBJECT_DIR):
            os.makedirs(self._config.OBJECT_DIR)
        error_paths: list[tuple[Path, str]] = []
        if self._config.JOB_COUNT > 1:  # Multi-process
            jobs: list[tuple[Path, subprocess.Popen]] = []
            completed_paths: list[Path] = []
            for source_path, cmd in self._todo_dict.items():
                print(ConsoleColor.BLUE + cmd + ConsoleColor.ENDCOLOR)
                jobs.insert(0, (source_path, subprocess.Popen(
                    cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True,
                    shell=True)))  # FIFO style (will be poped)
                if len(jobs) >= self._config.JOB_COUNT:  # If jobs count is maxed we wait for the oldest one to finished
                    job_path, oldest_job = jobs.pop()
                    oldest_job.wait()
                    if oldest_job.returncode != 0:
                        error_paths.append((job_path, oldest_job.stdout.read() + oldest_job.stderr.read()))
                        if str(job_path) in self._hash_dict:
                            del self._hash_dict[str(job_path)]
                        break
                    # Update hash if no error
                    for dependency_path in self._dependency_dict[str(job_path)]:
                        self._hash_dict[str(dependency_path)] = get_hash(dependency_path)
                    completed_paths.append(job_path)
            for source_path in completed_paths:
                del self._todo_dict[source_path]

            for job_path, job in jobs:  # Wait the last jobs to finish
                job.wait()
                if job.returncode != 0:
                    error_paths.append((job_path, job.stdout.read() + job.stderr.read()))
                    if str(job_path) in self._hash_dict:
                        del self._hash_dict[str(job_path)]
                else:
                    # Update hash if no error
                    for dependency_path in self._dependency_dict[str(job_path)]:
                        self._hash_dict[str(dependency_path)] = get_hash(dependency_path)
                    del self._todo_dict[job_path]
        else:  # Single-process
            completed_paths: list[Path] = []
            for source_path, cmd in self._todo_dict.items():
                print(ConsoleColor.BLUE + cmd + ConsoleColor.ENDCOLOR)
                job = subprocess.run(cmd, check=False, stdout=subprocess.PIPE, stderr=subprocess.PIPE,
                                     universal_newlines=True, shell=True)
                if job.returncode != 0:
                    error_paths.append((source_path, job.stdout + job.stderr))
                    if str(source_path) in self._hash_dict:
                        del self._hash_dict[str(source_path)]
                    break
                # Update hash if no error
                for dependency_path in self._dependency_dict[str(source_path)]:
                    self._hash_dict[str(dependency_path)] = get_hash(dependency_path)
                    completed_paths.append(source_path)
            for source_path in completed_paths:
                del self._todo_dict[source_path]

        if error_paths:
            for error_path, error_text in error_paths:
                print(ConsoleColor.RED + f'Error compiling {error_path}:\n' + ConsoleColor.ENDCOLOR + error_text)
            with open(self._config.OBJECT_DIR / 'hash.json', 'w', encoding='utf-8') as hash_file:
                hash_file.write(json.dumps(self._hash_dict, indent=1))
            return False

        # Running linking processes
        if not self._config.BIN_DIR.exists():
            self._config.BIN_DIR.mkdir(parents=True)
        all_app_objects = [self._config.OBJECT_DIR / Path(app_path).parent / (Path(app_path).stem + '.o')
                           for app_path in self._config.APPS]
        if self._config.JOB_COUNT > 1:  # Multi-process
            jobs: list[tuple[Path, subprocess.Popen]] = []
            for app_path, app_object_path in zip(self._config.APPS, all_app_objects):
                if app_path in self._config.IGNORE_APPS:
                    continue
                bin_path = self._config.BIN_DIR / app_path
                if not bin_path.parent.exists():
                    bin_path.parent.mkdir(parents=True)
                object_files = [str(object_path) for object_path in self._compile_dict.values()
                                if object_path not in all_app_objects]
                object_files.append(str(app_object_path))
                cmd = ' '.join([self._config.CC, self._config.COMMON_FLAGS, *object_files, '-o', str(bin_path),
                                self._config.LINK_FLAGS])
                print(ConsoleColor.BLUE + cmd + ConsoleColor.ENDCOLOR)
                jobs.insert(0, (app_path, subprocess.Popen(
                    cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True,
                    shell=True)))  # FIFO style (will be poped)
                if len(jobs) >= self._config.JOB_COUNT:  # If jobs count is maxed we wait for the oldest one to finished
                    app_path, oldest_job = jobs.pop()
                    oldest_job.wait()
                    if oldest_job.returncode != 0:
                        error_paths.append((app_path, oldest_job.stdout.read() + oldest_job.stderr.read()))
                        break
            for job_path, job in jobs:  # Wait the last jobs to finish
                job.wait()
                if job.returncode != 0:
                    error_paths.append((job_path, job.stdout.read() + job.stderr.read()))
        else:  # Single-process
            for app_path, app_object_path in zip(self._config.APPS, all_app_objects):
                if app_path in self._config.IGNORE_APPS:
                    continue
                bin_path = self._config.BIN_DIR / app_path
                if not bin_path.parent.exists():
                    bin_path.parent.mkdir(parents=True)
                object_files = [str(object_path) for object_path in self._compile_dict.values()
                                if object_path not in all_app_objects]
                object_files.append(str(app_object_path))
                cmd = ' '.join([self._config.CC, self._config.COMMON_FLAGS, *object_files, '-o', str(bin_path),
                                self._config.LINK_FLAGS])
                print(ConsoleColor.BLUE + cmd + ConsoleColor.ENDCOLOR)
                job = subprocess.run(cmd, check=False, stdout=subprocess.PIPE, stderr=subprocess.PIPE,
                                     universal_newlines=True, shell=True)
                if job.returncode != 0:
                    error_paths.append((app_path, job.stdout + job.stderr))
        if error_paths:
            for error_path, error_text in error_paths:
                print(ConsoleColor.RED + f'Error linking {error_path}:\n' + ConsoleColor.ENDCOLOR + error_text)
            with open(os.path.join(self._config.OBJECT_DIR, 'hash.json'), 'w', encoding='utf-8') as hash_file:
                hash_file.write(json.dumps(self._hash_dict, indent=1))
            return False

        with open(self._config.OBJECT_DIR / 'hash.json', 'w', encoding='utf-8') as hash_file:
            hash_file.write(json.dumps(self._hash_dict, indent=1))

        print(f'{ConsoleColor.GREEN}Compilation done{ConsoleColor.ENDCOLOR}')
        return True