#!/usr/bin/env python

import re
import glob
import sys
from collections import defaultdict
from typing import Dict, Set, List, Iterable


class Module:
    def __init__(self, name: str, depend: Iterable[str]) -> None:
        self.name = name
        self.depend = set(depend)
        self.depth = 0


class ModuleManager:
    """
    Class for resolving module dependencies.

    Attributes
    ----------
    modules: Dict[str, Module]
        registered modules
    """
    def __init__(self):
        self.modules = dict()

    def reset_depth(self) -> None:
        """
        Reset depth of all modules to 0.

        :return: None
        """
        for module in self.modules.values():
            module.depth = 0

    def check_sanity(self) -> None:
        """
        Check the sanity of modules.

        :return: None
        """
        for name, module in self.modules.items():
            # Check for cyclic dependencies
            if name in self.get_dependencies([name], include_roots=False):
                print(f"ERROR: module '{name}' has cyclic dependencies" )
                sys.exit(-1)

    def add_module(self, name: str, depend: Iterable[str]) -> None:
        """
        Add a new module

        :param name: name of the module
        :param depend: dependencies of the module
        :return: None
        """
        self.modules[name] = Module(name, depend)

    def get_dependencies(self, root_mods: Iterable[str],
                         include_roots: bool = True) -> Set[str]:
        """
        Extract the dependencies recursively for a list of modules.

        :param root_mods: names of root modules
        :param include_roots: whether to include root mods in results
        :return: set of the names of all the dependencies
        """
        if include_roots:
            dependencies = [_ for _ in root_mods]
        else:
            dependencies = []
            for mod_name in root_mods:
                dependencies.extend(self.modules[mod_name].depend)
        for mod_name in dependencies:
            module = self.modules[mod_name]
            for depend_item in module.depend:
                if depend_item not in dependencies:
                    dependencies.append(depend_item)
        return set(dependencies)

    def sort_modules(self, mod_list: Iterable[str]) -> List[Module]:
        """
        Sort modules according to their depth in the dependency tree.

        :param mod_list: names of modules to sort
        :return: sorted modules with depth in increasing order
        """
        self.reset_depth()
        depend_tree = [self.modules[_] for _ in mod_list]
        depth = 0
        updated = True
        while updated:
            updated = False
            nodes_to_check = [_ for _ in depend_tree if _.depth == depth]
            for parent in nodes_to_check:
                for child in depend_tree:
                    if child.name in self.modules[parent.name].depend:
                        child.depth = depth + 1
                        updated = True
            depth += 1
            if depth >= 255:
                raise RuntimeError("Possible cyclic dependency")
        sorted_modules = sorted(depend_tree, key=lambda x: (x.depth, x.name))
        return sorted_modules


def find_usages() -> Dict[str, Set[str]]:
    """
    Find the modules referenced by each module under 'sources'.

    :return: dictionary with keys being module names and values being the
        dependencies
    """
    # source_files = ['sources/atoms/atoms.h', 'sources/base/parser.h', ...]
    # usages = {'atoms': {'atoms.h', 'vector', 'pseudo_pot.h', 'eigen3', 'string'},
    #           'dynamics': {'vector', 'eigen3', 'dynamics.h', 'string'},}
    source_files = glob.glob("sources/*/*.h")
    source_files.extend(glob.glob("sources/*/*.cpp"))
    source_files.extend(glob.glob("sources/*/*.c"))
    usages = defaultdict(set)
    pattern = re.compile(r"^\s*#\s*include\s*[<\"]+\s*([\w/.]+)[>\"]+\s*", re.I)
    for src in source_files:
        module = src.split("/")[1]
        with open(src, "r") as src_file:
            content = src_file.readlines()
        for line in content:
            result = re.search(pattern, line)
            if re.search(pattern, line) is not None:
                usages[module].add(result.group(1).split("/")[0])
    return usages


def print_depend_all(mod_man: ModuleManager) -> None:
    """
    Print dependencies for sources/CMakeLists.txt.

    :param mod_man: module manager containing the dependency info
    """
    print_cuda = lambda x: print(f"if(WITH_CUDA)\n{'':4s}{x}\nendif()")
    sorted_modules = sorted(mod_man.modules.keys())
    for name in sorted_modules:
        depend = sorted(mod_man.get_dependencies([name], include_roots=False))
        if len(depend) > 0:
            depend_str = [f" {_}" for _ in depend if _ != "cuda"]  # filter cuda
            depend_str = "".join(depend_str).lstrip()
            if name == "cuda":
                scope = "${CUDA_LINK_LIBRARIES_KEYWORD}"
                print_cuda(f"target_link_libraries({name} {scope} {depend_str})")
            else:
                scope = "PUBLIC"
                print(f"target_link_libraries({name} {scope} {depend_str})")
                if "cuda" in depend:
                    print_cuda(f"target_link_libraries({name} {scope} cuda)")
            print()


def print_depend_each(mod_man: ModuleManager) -> None:
    """
    Print dependencies for sources/xxx/CMakeLists.txt.

    :param mod_man: module manager containing the dependency info
    """
    print_cuda = lambda x: print(f"if(WITH_CUDA)\n{'':4s}{x}\nendif()")
    sorted_modules = sorted(mod_man.modules.keys())
    for name in sorted_modules:
        depend = sorted(mod_man.get_dependencies([name], include_roots=False))
        if len(depend) > 0:
            depend_str = [f" {_}" for _ in depend if _ != "cuda"]  # filter cuda
            depend_str = "".join(depend_str).lstrip()
            print(f"\n--------- CMake settings for {name}  --------\n")
            if name == "cuda":
                scope = "${CUDA_LINK_LIBRARIES_KEYWORD}"
                print_cuda(f"target_link_libraries(${{target_name}} {scope} {depend_str})")
            else:
                scope = "PUBLIC"
                print(f"target_link_libraries(${{target_name}} {scope} {depend_str})")
                if "cuda" in depend:
                    print_cuda(f"target_link_libraries(${{target_name}} {scope} cuda)")
            print()


def write_cmake_config(name: str, config: List[str]) -> None:
    """
    Write dependencies into sources/name/CMakeLists.txt.

    :param name: subdirectory name
    :param config: dependencies
    """
    with open(f"sources/{name}/CMakeLists.txt", "r") as cmake_file:
        original_content = cmake_file.readlines()

    new_content = []
    in_block = False  # Whether in 'link_order' block
    for line in original_content:
        if line.find("begin_link_order") != -1:
            in_block = True
            new_content.append(line)
            new_content.extend(config)
        if line.find("end_link_order") != -1:
            in_block = False
        if not in_block:
            new_content.append(line)

    with open(f"sources/{name}/CMakeLists.txt", "w") as cmake_file:
        cmake_file.writelines(new_content)


def write_depend_each(mod_man: ModuleManager) -> None:
    """
    Write dependencies into sources/xxx/CMakeLists.txt.

    :param mod_man: module manager containing the dependency info
    """
    print_cuda = lambda x: f"if(WITH_CUDA)\n{'':4s}{x}\nendif()"
    sorted_modules = sorted(mod_man.modules.keys())
    for name in sorted_modules:
        depend = sorted(mod_man.get_dependencies([name], include_roots=False))
        if len(depend) > 0:
            depend_str = [f" {_}" for _ in depend if _ != "cuda"]  # filter cuda
            depend_str = "".join(depend_str).lstrip()
            config = []
            if name == "cuda":
                scope = "${CUDA_LINK_LIBRARIES_KEYWORD}"
                config.append(print_cuda(f"target_link_libraries(${{target_name}} {scope} {depend_str})"))
            else:
                scope = "PUBLIC"
                config.append(f"target_link_libraries(${{target_name}} {scope} {depend_str})")
                if "cuda" in depend:
                     config.append(print_cuda(f"target_link_libraries(${{target_name}} {scope} cuda)"))

            # Append \n to line ends
            config = [f"{_}\n" for _ in config]
            write_cmake_config(name, config)


def resolve_depend() -> None:
    """
    Resolve dependencies between modules.

    :return: None
    """
    usages = find_usages()
    modules = set(usages.keys())

    # Filter external and self-dependencies
    mod_man = ModuleManager()
    for name, depend in usages.items():
        intra_depend = modules.intersection(depend).difference({name})
        mod_man.add_module(name, intra_depend)
    mod_man.check_sanity()

    # print_depend_all(mod_man)
    # print_depend_each(mod_man)
    write_depend_each(mod_man)


def main():
    resolve_depend()


if __name__ == "__main__":
    main()
