from IPython.core.magic import register_cell_magic, needs_local_scope

import importlib
import glob
import sys, os
import random
import string
import hashlib
import shutil
import subprocess
import shlex

from lief import parse, ELF

compiler_cmd = "c++"

optimize = "-O0"

cmd_flags = ["$(pkg-config --cflags eigen3)", "-Wall"]

template = """
#include <pybind11/pybind11.h>
namespace py = pybind11;

{cell_code}

PYBIND11_MODULE({name}, m)
    {defs}
"""

def_template = '\nm.def("{function}", &{function});\n'

TMP_PATH = "pbm_modules"

if os.path.exists(TMP_PATH):
    shutil.rmtree(TMP_PATH)

BUILD_DIR = os.path.join(TMP_PATH, "build")
MODULES_DIR = os.path.join(TMP_PATH, "modules")
LIB_DIR = os.path.join(TMP_PATH, "lib")
# Make sure the correct python executeable is used using cmake
os.environ["PATH"] = os.path.dirname(sys.executable) + os.pathsep + os.environ["PATH"]

# def run_command(command):
#     process = subprocess.Popen(shlex.split(command), stdout=subprocess.PIPE)
#     while True:
#         output = process.stdout.readline()
#         if output == '' and process.poll() is not None:
#             break
#         if output:
#             print output.strip()
#     rc = process.poll()
#     return rc

def exec_command(cmd, cwd=None,**kwargs):
    if isinstance(cmd, str):
        cmd = shlex.split(cmd)
    try:
        output = subprocess.check_output(cmd,cwd=cwd,stderr=subprocess.STDOUT,**kwargs)
        if isinstance(output,bytes):
            return False, output.decode("utf-8")
        return False, output
    except subprocess.CalledProcessError as e:
        if isinstance(e.output,bytes):
            return True, e.output.decode("utf-8")
        return True, e.output


def ensure_tmp_folder(path=TMP_PATH):
    path = os.path.join(os.getcwd(), path)
    if not os.path.exists(path):
        os.makedirs(path)


def ensure_pybind(path, force=False):
    path = os.path.join(path,"pybind11")
    exists = os.path.exists(path)
    if exists and force:
        shutil.rmtree(os.path.join(path,"pybind11"))
    if not exists:
        err, output = exec_command(["git", "clone", "https://github.com/pybind/pybind11.git", path])


def ensure_build_dir(path, clear=False):
    path = os.path.join(path, "build")
    exists = os.path.exists(path)
    if exists and clear:
        shutil.rmtree(path)
    if not exists or clear:
        os.makedirs(path)
    return exists and not clear


def ensure_linker():
    path = os.path.join(TMP_PATH, "liblinker.so")
    exists = os.path.exists(path)
    if not os.path.exists(path):
        return compile_linkerlib()
    return


def parse_args(line):
    args = line.split(" ")
    ret = {}
    ret["cmake_path"] = None
    ret["functions"] = []
    ret["rebuild"] = False
    # TODO: Add full argument parser
    for i,arg in enumerate(args):
        if arg == "-c":
            ret["cmake_path"] = args[i+1].strip()
        if arg == "-f":
            ret["functions"].append(args[i+1].strip())
        if arg == "-rebuild":
            ret["rebuild"] = True
    return ret


def get_CMakeLists(args):
    CMakeLists = ""
    if args["cmake_path"] is not None:
        with open(os.path.join(args["cmake_path"],"CMakeLists.txt"), 'r') as f:
            CMakeLists = f.read()
    return CMakeLists


def cleanup(args, name, CMakeLists):
    if args["cmake_path"] is not None:
        with open(os.path.join(args["cmake_path"],"CMakeLists.txt"), 'w') as f:
            f.write(CMakeLists)
    if os.path.exists(os.path.join(args["path"],f"{name}.cpp")):
        os.remove(os.path.join(args["path"],f"{name}.cpp"))


def get_gcc_cmd(name):
    pyexe = sys.executable
    return f"{compiler_cmd} {' '.join(cmd_flags)} {optimize} -std=c++11 -fPIC $({pyexe} -m pybind11 --includes) -c {BUILD_DIR}/{name}.cpp -o {BUILD_DIR}/{name}.o"


def get_linker_cmd(name, others=None):
    if others is None:
        others = " ".join([f"-l{os.path.basename(n)[3:-3]}" for n in sorted(glob.glob(f"{LIB_DIR}/*.so"), key=os.path.getmtime)[::-1]])
    return f"{compiler_cmd} -Wl,-rpath {LIB_DIR} -shared {BUILD_DIR}/{name}.o -o {MODULES_DIR}/{name}.so -L{LIB_DIR} {others}"


def import_to_ip(name):
    # get a handle on the module
    mdl = importlib.import_module(name)
    # is there an __all__?  if so respect it
    if "__all__" in mdl.__dict__:
        names = mdl.__dict__["__all__"]
    else:
        # otherwise we import all names that don't begin with _
        names = [x for x in mdl.__dict__ if not x.startswith("_")]
    get_ipython().push({k: getattr(mdl, k) for k in names})


def get_importexport(file):
    binary = parse(file)
    symbols = binary.symbols
    imported = set()
    exported = set()
    for i,symbol in enumerate(symbols):
        if not symbol.imported and not symbol.exported and symbol.binding == ELF.SYMBOL_BINDINGS.GLOBAL:
            imported.add(symbol.name)
        if symbol.exported and symbol.binding == ELF.SYMBOL_BINDINGS.GLOBAL:
            exported.add(symbol.name)
    return imported, exported


export_graph = {}
import_graph = {}
symbols = {}

def update_dependencies(fresh_compiled):
    global export_graph, import_graph
    imports, exports = get_importexport(fresh_compiled)
    symbols[fresh_compiled] = (imports, exports)
    for inp in imports:
        if inp.startswith("Py") or inp.startswith("_Py"):
            continue
        if inp in import_graph:
            import_graph[inp].append(fresh_compiled)
        else:
            import_graph[inp] = [fresh_compiled]

    for ex in exports:
        if ex.startswith("PyInit"):
            continue
        export_graph[ex] = fresh_compiled


def get_list_of_rebuilds(fresh_compiled, need_recompilation=None):
    global export_graph, import_graph
    if need_recompilation is None:
        need_recompilation = set()
        need_recompilation.add(fresh_compiled)
    exports = symbols[fresh_compiled][1]
    for ex in exports:
        if ex in import_graph:
            for lib in import_graph[ex]:
                if lib not in need_recompilation:
                    need_recompilation.add(lib)
                    # recursive collect all libs, that depend on this
                    get_list_of_rebuilds(lib, need_recompilation)
    return need_recompilation

def get_list_of_links(fresh_compiled):
    global export_graph, import_graph
    imports = symbols[fresh_compiled][0]
    res = []
    for imp in imports:
        if imp in export_graph:
            lib = export_graph[imp]
            if lib not in res:
                res.append(lib)
    return res


def bulk_import(need_reimport=None):
    if need_reimport is None:
        need_reimport = sorted(glob.glob(f"{MODULES_DIR}/*.so"), key=os.path.getmtime)

    # create unique folder
    rnd = ''.join(random.choice(string.ascii_uppercase + string.digits) for _ in range(5))
    rnd_path = os.path.join(TMP_PATH, rnd)
    if not os.path.exists(rnd_path):
        os.makedirs(rnd_path)
        with open(os.path.join(rnd_path,"__init__.py"), 'w+') as f:
            pass

    for lib_path in need_reimport:
        lib = os.path.basename(lib_path)[:-3]
        shutil.copyfile(lib_path, os.path.join(rnd_path,f"{lib}.so"))

    for lib_path in need_reimport:
        lib = os.path.basename(lib_path)[:-3]
        import_to_ip(f"{rnd}.{lib}")
    shutil.rmtree(rnd_path)


def compile_module(cell, args):
    name = args["name"]
    ensure_tmp_folder(BUILD_DIR)

    create_cpp_file(BUILD_DIR, cell, args)

    ip = get_ipython()
    # We compile the C++ code into an object file
    command = get_gcc_cmd(name)
    compile = ip.getoutput(command)
    # currently warnings will lead to an abort
    if len(compile) != 0:
        raise Exception("\n".join(compile))

    # link into a shared object, without links
    linker = get_linker_cmd(name, "")
    link = ip.getoutput(linker)
    if len(link) != 0:
        raise Exception("\n".join(link))


def link_module(name, others=None):
    other_name = os.path.basename(name)[:-3]
    linker = get_linker_cmd(other_name, others)
    link = get_ipython().getoutput(linker)
    if len(link) != 0:
        raise Exception("\n".join(link))


version = 0

def create_cpp_file(path, cell, args):
    name = args["name"]
    # add pybind11 function defs
    split = cell.split("defs")
    if len(split) == 1 and len(args["functions"]) == 0:
        raise Exception("You have to name a function as argument ('%%cpp -f <functionname>') or manual set defs")
    # make sure that split is at least length 2
    split.append("")
    curr_defs = split[1]
    for function in args["functions"]:
        curr_defs += def_template.format(function=function)
    with open(os.path.join(path,f"{name}.cpp"), 'w') as f:
        f.write(template.format(cell_code=split[0],name=name,defs="{"+curr_defs+"}"))


def get_newest(name):
    return sorted(glob.glob(f"{LIB_DIR}/lib{name}*.so"), key=os.path.getmtime)[-1]

history = []


# TODO: TEST
@register_cell_magic
def cpp(line, cell, *what):
    global version
    args = parse_args(line)
    ip = get_ipython()
    # get name of module
    if args["cmake_path"] is None:
        # TODO!
        compile_id = get_gcc_cmd("dummy")
    else:
        compile_id = get_CMakeLists(args)

    hash_object = hashlib.sha256((compile_id+"|"+line+"|"+cell).encode())
    name = "cpp_magic_"+hash_object.hexdigest()[:10]
    args["name"] = name

    version += 1
    history.append(os.path.join(MODULES_DIR, f"{name}.so"))

    ensure_tmp_folder(TMP_PATH)
    # Check if neccessary here
    if not TMP_PATH in sys.path:
        sys.path.append(TMP_PATH)
    ensure_tmp_folder(MODULES_DIR)
    ensure_tmp_folder(LIB_DIR)

    # if not yet compiled, do so
    if not os.path.exists(os.path.join(MODULES_DIR,f"{name}.so")):
        # TODO: extend to cmake building
        compile_module(cell, args)

    update_dependencies(os.path.join(MODULES_DIR,f"{name}.so"))
    relinks = list(get_list_of_rebuilds(os.path.join(MODULES_DIR,f"{name}.so")))
    relinks = sorted(relinks, key=os.path.getmtime)[::-1]
    for re in relinks:
        other_name = os.path.basename(re)[:-3]
        shutil.copyfile(os.path.join(MODULES_DIR,f"{other_name}.so"), os.path.join(LIB_DIR,f"lib{other_name}{version}.so"))
        # touch
        open(os.path.join(LIB_DIR,f"lib{other_name}{version}.so"), 'a').close()
        # TODO: delete old


    # modules = sorted(glob.glob(f"{MODULES_DIR}/*.so"), key=os.path.getmtime)[::-1]
    # others = " ".join([f"-l{os.path.basename(get_newest(os.path.basename(n)[:-3]))[3:-3]}" for n in modules])
    # print(others)
    for re in relinks:
        others = " ".join([f"-l{os.path.basename(get_newest(os.path.basename(n)[:-3]))[3:-3]}" for n in sorted(get_list_of_links(re), key=os.path.getmtime)[::-1]])
        link_module(re, others)
        other_name = os.path.basename(re)[:-3]
        shutil.copyfile(os.path.join(MODULES_DIR,f"{other_name}.so"), os.path.join(LIB_DIR,f"lib{other_name}{version}.so"))

    bulk_import([h for h in history if h in relinks])

"""
CLEAR LIB_FOLDER
#### PIPELINE ####

create name
add version
add name to history
compile to modules_dir/{name}.so (may be already there)
detect_relink_needs

for all relinked mv modules_dir/{name}.so -> libs_dir/lib{name}{version}.so (delete old?)
relink
for all relinked mv modules_dir/{name}.so -> libs_dir/lib{name}{version}.so
import_bulk according to history

"""