#! python3

from argparse import ArgumentParser
from dataclasses import dataclass, field
import hashlib
import json
import os
from pathlib import Path
import subprocess
import shutil
import sys
from typing import Any, Callable, Iterable


@dataclass
class Config:
    cc: str = 'g++'                      # Compiler to call
    # Output binaries (need to be found as .cpp directly in source_dir)
    apps: Iterable[str] = field(default_factory=lambda: ['app_name'])
    ignore_apps: Iterable[str] = field(default_factory=list)
    job_count: int = int((os.cpu_count() or 1) * 0.8)  # Concurent jobs (multi-processing)
    watch: bool = False                  # Watch source modification for auto-compiling

    bin_dir: Path = Path('bin')          # Output directory (binaries)
    include_dir: Path = Path('include')  # Include directory (header files)
    object_dir: Path = Path('obj')       # Temporary directory (object files)
    source_dir: Path = Path('src')       # Source directories

    common_flags: str = '-std=c++17'               # Flags used for comiling and linking
    common_debug_flags: str = '-g'                 # Flags added in debug mode
    common_release_flags: str = '-O2 -flto'        # Flags added in release mode
    compile_flags: str = f'-Wall -I{include_dir}'  # Flags added for compiling (recommandation : `pkg-config --cflags`)
    link_flags: str = ''                           # Flags added for linking (recommandation : `pkg-config --libs`)

    cpp_sources: Iterable[Path] = field(default_factory=lambda: [
        filepath for filepath in Path('src').rglob('*.cpp') if not filepath.name.startswith('.')])

    # Function to run before compile (not call in --clean situation)
    pre_compile_function: Callable[[], Any] | None = None


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(hash_path: Path) -> str:
    hash_obj = hashlib.md5()
    hash_obj.update(hash_path.read_text(encoding='utf-8').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='release', 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()

    # Update flags and directories for mode debug/release
    if arguments.type == 'debug':
        config.common_flags += ' ' + config.common_debug_flags
        config.object_dir /= 'debug'
    else:
        config.common_flags += ' ' + config.common_release_flags
        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[Path, str] = {}
        hash_dict_path = self._config.object_dir / 'hash.json'
        if hash_dict_path.exists():
            try:
                self._hash_dict = {
                    Path(path): value for path, value in json.loads(hash_dict_path.read_text(encoding='utf-8')).items()}
            except json.JSONDecodeError:
                pass

        # Get source dependencies
        error_paths: list[tuple[Path, str]] = []
        self._dependency_dict: dict[Path, list[Path]] = {}
        for source_path in self._compile_dict:
            cmd = ' '.join(
                [self._config.cc, self._config.common_flags, self._config.compile_flags, str(source_path), '-M'])
            job = subprocess.run(cmd, check=False, capture_output=True, text=True, shell=True)
            if job.returncode != 0:
                error_paths.append((source_path, job.stdout + job.stderr))
                break
            self._dependency_dict[source_path] = [
                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}:{ConsoleColor.ENDCOLOR}'
                      f'\n{error_text}')
            sys.exit(1)

    def make(self):
        self._populate_todo_dict()

        if self._config.watch:
            from watch import Watcher, WatchFlag  # ruff:ignore[import-outside-top-level]

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

            self._compile_sources()
            print(f'{ConsoleColor.GREEN}Watching for changes{ConsoleColor.ENDCOLOR}')
            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 _save_hash_dict(self):
        (self._config.object_dir / 'hash.json').write_text(
            json.dumps({str(path): value for path, value in self._hash_dict.items()}, indent=1), encoding='utf-8')

    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[source_path]:
                dependency_hash = get_hash(path)
                if path not in self._hash_dict or self._hash_dict[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 source_path not in self._hash_dict or self._hash_dict[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 self._config.object_dir.exists:
            self._config.object_dir.mkdir(parents=True)
        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(f'{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:
                        if oldest_job.stdout is None or oldest_job.stderr is None:
                            raise RuntimeError('Unexpected error: job without stdout/stdin')
                        error_paths.append((job_path, oldest_job.stdout.read() + oldest_job.stderr.read()))
                        if job_path in self._hash_dict:
                            del self._hash_dict[job_path]
                        break
                    # Update hash if no error
                    for dependency_path in self._dependency_dict[job_path]:
                        self._hash_dict[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:
                    if job.stdout is None or job.stderr is None:
                        raise RuntimeError('Unexpected error: job without stdout/stdin')
                    error_paths.append((job_path, job.stdout.read() + job.stderr.read()))
                    if job_path in self._hash_dict:
                        del self._hash_dict[job_path]
                else:
                    # Update hash if no error
                    for dependency_path in self._dependency_dict[job_path]:
                        self._hash_dict[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(f'{ConsoleColor.BLUE}{cmd}{ConsoleColor.ENDCOLOR}')
                job = subprocess.run(cmd, check=False, capture_output=True, text=True, shell=True)
                if job.returncode != 0:
                    error_paths.append((source_path, job.stdout + job.stderr))
                    if source_path in self._hash_dict:
                        del self._hash_dict[source_path]
                    break
                # Update hash if no error
                for dependency_path in self._dependency_dict[source_path]:
                    self._hash_dict[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(f'{ConsoleColor.RED}Error compiling {error_path}:{ConsoleColor.ENDCOLOR}\n{error_text}')
            self._save_hash_dict()
            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: list[str] = [str(app_object_path), *[
                    str(object_path) for object_path in self._compile_dict.values()
                    if object_path not in all_app_objects]]
                cmd = ' '.join([self._config.cc, self._config.common_flags, *object_files, '-o', str(bin_path),
                                self._config.link_flags])
                print(f'{ConsoleColor.BLUE}{cmd}{ConsoleColor.ENDCOLOR}')
                jobs.insert(0, (Path(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:
                        if oldest_job.stdout is None or oldest_job.stderr is None:
                            raise RuntimeError('Unexpected error: job without stdout/stdin')
                        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:
                    if job.stdout is None or job.stderr is None:
                        raise RuntimeError('Unexpected error: job without stdout/stdin')
                    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: list[str] = [str(app_object_path), *[
                    str(object_path) for object_path in self._compile_dict.values()
                    if object_path not in all_app_objects]]
                cmd = ' '.join([self._config.cc, self._config.common_flags, *object_files, '-o', str(bin_path),
                                self._config.link_flags])
                print(f'{ConsoleColor.BLUE}{cmd}{ConsoleColor.ENDCOLOR}')
                job = subprocess.run(cmd, check=False, capture_output=True, text=True, shell=True)
                if job.returncode != 0:
                    error_paths.append((Path(app_path), job.stdout + job.stderr))
        self._save_hash_dict()
        if error_paths:
            for error_path, error_text in error_paths:
                print(f'{ConsoleColor.RED}Error linking {error_path}:{ConsoleColor.ENDCOLOR}\n{error_text}')
            return False

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