import importlib
import inspect


LIBS = {
    "base": [
        {
            "os": "os",
            "glob": "glob",
            "math": "math",
            "time": "time",
            "random": "random",
            "argparse": "argparse",
        },
        (),
    ],
    "utils": [
        {
            "tqdm": "tqdm",
            "rich": "rich",
        },
        ("base"),
    ],  # dependency
    "data": [
        {
            "np": "numpy",  # alternatives
            "cv2": "cv2",
            "plt": "matplotlib.pyplot",
            "Image": "PIL.Image",
            # 'vis': vis, # live object
        },
        ("base", "utils"),
    ],
    "torch": [
        {
            "torch": "torch",
            "nn": "torch.nn",
            "F": "torch.nn.functional",
            "Dataset": ("torch.utils.data", "Dataset"),  # single class/function
            "DataLoader": ("torch.utils.data", "DataLoader"),
        },
        ("base", "utils", "data"),
    ],
}

LIBS["all"] = [{}, tuple(LIBS.keys())]

G = None


def retrieve_globals(verbose=False):
    # locate and set G to the globals which directly `import kiui`. (only once)
    # ref: https://stackoverflow.com/questions/40652688/how-to-access-globals-of-parent-module-into-a-sub-module/50381748
    global G
    stack = inspect.stack()
    frame_id = 1
    while frame_id < len(stack):
        g = dict(inspect.getmembers(stack[frame_id][0]))["f_globals"]
        if "kiui" in g:
            G = g
            if verbose:
                print(f"[KiuiKit-INFO] located global frame at {frame_id}")
            break
        frame_id += 1
    if G is None:
        raise RuntimeError(
            "Cannot locate global frame, make sure you called exactly `import kiui`!"
        )


def try_import(target, sources, verbose=False):

    if G is None:
        retrieve_globals(verbose)

    if target in G:
        if verbose:
            print(f"[KiuiKit-INFO] {target} is already present, skipped.")
        return

    if not isinstance(sources, list):
        sources = [sources]

    for source in sources:
        try:
            if verbose:
                print(f"[KiuiKit-INFO] try to import {source}")

            # (module, component) or ("module", component)
            if isinstance(source, tuple):
                source_module, source_component = source
                if isinstance(source_module, str):
                    source_module = importlib.import_module(source_module)
                source = getattr(source_module, source_component)
            # "module"
            elif isinstance(source, str):
                source = importlib.import_module(source)
            # module
            else:
                pass

            G[target] = source

            if verbose:
                print(f"[KiuiKit-INFO] succeed to import {source} as {target}")
            break

        except ImportError as e:
            print(f"[KiuiKit-WARN] failed to import {source} as {target}: {str(e)}")


def import_libs(pack, verbose=False):
    for k, v in LIBS[pack][0].items():
        try_import(k, v, verbose)


"""
setup all import in one line.
usage:
    kiui.env() # setup base env
    kiui.env('torch') # pytorch with all regular dependencies
    kiui.env('[torch]') # pytorch only
"""


def env(*packs, verbose=False):
    if len(packs) == 0:
        packs = ["base"]

    def check_pack(pack):
        if pack not in LIBS:
            raise RuntimeError(
                f"[Kiui-ERROR] unknown pack {pack}, availables: {list(LIBS.keys())}"
            )

    def resolve_env(packs):
        res = set()
        for pack in packs:
            if pack[0] == "[" and pack[-1] == "]":
                pack = pack[1:-1]
                check_pack(pack)
                res.add(pack)
            else:
                check_pack(pack)
                for dep in LIBS[pack][1]:
                    res.add(dep)
                res.add(pack)
        return list(res)

    packs = resolve_env(packs)
    for pack in packs:
        import_libs(pack, verbose)
