from importlib.util import module_from_spec, spec_from_file_location
from os import getcwd, path, walk
from pathlib import Path
from typing import Union

# =================================================================================


# Import functions
# Took from https://stackoverflow.com/a/57892961
def get_py_files(src):
    cwd = getcwd()  # Current Working directory
    py_files = []
    for root, dirs, files in walk(src):
        py_files.extend(
            Path(f"{cwd}/{root}/{file}") for file in files if file.endswith(".py")
        )
    return py_files


def dynamic_import(module_name: str, py_path: str):
    try:
        module_spec = spec_from_file_location(module_name, py_path)
        if module_spec is None:
            raise RuntimeError(
                f"Module spec from module name {module_name} and path {py_path} is None"
            )
        module = module_from_spec(module_spec)
        module_spec.loader.exec_module(module)
        return module
    except SyntaxError:
        print(
            f"Could not load extension {module_name} due to invalid syntax. Check logs/errors.log for details.",
            flush=True,
        )
        return
    except Exception as exc:
        print(f"Could not load extension {module_name} due to {exc}", flush=True)
        return


def dynamic_import_from_src(src: Union[str, Path], star_import=False):
    my_py_files = get_py_files(src)
    for py_file in my_py_files:
        module_name = Path(py_file).stem
        print(f"Importing {module_name} extension...", flush=True)
        imported_module = dynamic_import(module_name, py_file)
        if imported_module != None:
            if star_import:
                for obj in dir(imported_module):
                    globals()[obj] = imported_module.__dict__[obj]
            else:
                globals()[module_name] = imported_module
        print(f"Successfully loaded {module_name} extension", flush=True)
    return


# =================================================================================