import glob
import json
import inspect
import os
from os.path import (
    abspath,
    basename,
    dirname,
    exists,
    isdir,
    join,
    normpath,
    relpath,
    sep,
    splitext,
)
import pathlib
import posixpath
import shutil
import sysconfig
import toposort
from typing import Any, Dict, List, Optional, Set, Tuple

from urllib.error import HTTPError
import dataclasses

from setuptools import Extension

from cxxheaderparser.options import ParserOptions
from cxxheaderparser import preprocessor


from .download import download_and_extract_zip
from .config.pyproject_toml import PatchInfo, WrapperConfig, Download

from .autowrap.cxxparser import parse_header
from .autowrap.generator_data import GeneratorData, MissingReporter
from .autowrap.writer import WrapperWriter

from .config.autowrap_yml import AutowrapConfigYaml
from .config.dev_yml import get_dev_config
from .config.pyproject_toml import WrapperConfig, Download

# TODO: eventually provide native preprocessor by default and allow it
#       to be enabled/disabled per-file just in case
if os.getenv("RPYBUILD_PP_GCC") == "1":
    # GCC preprocessor can be 10x faster than pcpp for very complex files
    def make_preprocessor(*args, **kwargs):
        return preprocessor.make_gcc_preprocessor(print_cmd=False, *args, **kwargs)

else:
    make_preprocessor = preprocessor.make_pcpp_preprocessor


class Wrapper:
    """
    Wraps downloading bindings and generating them
    """

    # Used during preprocessing
    # -> should we change this based on what flags the compiler supports?
    _cpp_version = "__cplusplus 201703L"

    def __init__(self, package_name, cfg: WrapperConfig, setup, wwriter: WrapperWriter):
        self.package_name = package_name
        self.cfg = cfg
        self.wwriter = wwriter

        self.setup_root = setup.root
        self.pypi_package = setup.pypi_package
        self.root = join(setup.root, *package_name.split("."))

        # must match PkgCfg.name
        self.name = cfg.name
        self.static_lib = False

        # Compute the extension name, even if we don't create one
        extname = cfg.extension
        if not extname:
            extname = f"_{cfg.name}"

        # must match PkgCfg.libinit_import
        if cfg.libinit:
            libinit_py = cfg.libinit
            if libinit_py == "__init__.py":
                self.libinit_import = package_name
            else:
                pkg = splitext(libinit_py)[0]
                self.libinit_import = f"{package_name}.{pkg}"
        else:
            libinit_py = f"_init{extname}.py"
            self.libinit_import = f"{package_name}._init{extname}"

        self.libinit_import_py = join(self.root, libinit_py)

        self.platform = setup.platform
        self.pkgcfg = setup.pkgcfg

        # Used by pkgcfg
        self.depends = self.cfg.depends

        # Files that are generated AND need to be in the final wheel. Used by build_py
        self.additional_data_files: List[str] = []

        self._all_deps = None

        self._gen_includes = []

        self.extension = None
        if self.cfg.sources or self.cfg.autogen_headers:
            define_macros = [("RPYBUILD_MODULE_NAME", extname)] + [
                tuple(d.split(" ")) for d in self.platform.defines
            ]
            define_macros += [tuple(m.split(" ", 1)) for m in self.cfg.pp_defines]

            # extensions just hold data about what to actually build, we can
            # actually modify extensions all the way up until the build
            # really happens
            extname_full = f"{self.package_name}.{extname}"
            self.extension = Extension(
                extname_full,
                self.cfg.sources,
                define_macros=define_macros,
                language="c++",
            )

            # Add self to extension so that build_ext can query it on OSX
            self.extension.rpybuild_wrapper = self

            # Used if the maven download fails
            self.supported_platforms = setup.project.supported_platforms

        if self.cfg.autogen_headers and not self.cfg.generation_data:
            raise ValueError(
                "generation_data must be specified when autogen_headers/generate is specified"
            )

        # Setup an entry point (written during build_clib)
        entry_point = f"{self.cfg.name} = {self.package_name}.pkgcfg"

        setup_kwargs = setup.setup_kwargs
        ep = setup_kwargs.setdefault("entry_points", {})
        ep.setdefault("robotpybuild", []).append(entry_point)

        self.incdir = join(self.root, "include")
        self.rpy_incdir = join(self.root, "rpy-include")

        self.dev_config = get_dev_config(self.name)

        self._update_addl_data_files()

    def _extract_zip_to(self, dl: Download, dst, cache):
        try:
            download_and_extract_zip(dl.url, dst, cache)
        except HTTPError as e:
            # Check for a 404 error and raise an error if the platform isn't supported.
            if e.code != 404:
                raise e
            else:
                platform_dict = dataclasses.asdict(self.platform)

                os = platform_dict["os"]
                arch = platform_dict["arch"]

                is_os_supported = False
                is_arch_supported = False

                for supp_plat in self.supported_platforms:
                    if supp_plat.os is None or supp_plat.os == os:
                        is_os_supported = True
                        if supp_plat.arch is None or supp_plat.arch == arch:
                            is_arch_supported = True

                if not (is_os_supported and is_arch_supported):
                    if arch == "x86":
                        arch = "32-bit"
                    elif arch == "x86-64":
                        arch = "64-bit"

                    if os == "osx":
                        os = "macOS"

                    if not is_os_supported:
                        arch = ""

                    msg_plat = "{}{}{}".format(arch, " " if arch != "" else "", os)

                    err_msg = "{} is not supported on {}!".format(
                        self.pypi_package, msg_plat
                    )

                    raise OSError(err_msg)
                raise e

    def _add_addl_data_file(self, fullpath):
        if not isdir(fullpath):
            self.additional_data_files.append(relpath(fullpath, self.root))

    # pkgcfg interface
    def get_include_dirs(self) -> List[str]:
        includes = [self.incdir, self.rpy_incdir]
        if self.cfg.download:
            for dl in self.cfg.download:
                if dl.extra_includes:
                    includes += [join(self.incdir, inc) for inc in dl.extra_includes]
        for h in self.cfg.extra_includes:
            includes.append(join(self.setup_root, normpath(h)))
        return includes

    def get_library_dirs(self) -> Optional[List[str]]:
        if self.get_library_full_names():
            return [join(self.root, "lib")]
        return []

    def get_library_dirs_rel(self) -> Optional[List[str]]:
        if self.get_library_full_names():
            return ["lib"]
        return []

    def get_library_names(self) -> Optional[List[str]]:
        libs = []
        if self.cfg.download:
            for dl in self.cfg.download:
                if dl.libs:
                    libs += dl.libs
        return libs

    def get_library_full_names(self) -> Optional[List[str]]:
        if not self.cfg.download:
            return []

        dlopen_libnames = self.get_dlopen_library_names()

        libnames_full = []

        for dl in self.cfg.download:
            libext = dl.libexts.get(self.platform.libext, self.platform.libext)
            if dl.libs:
                for lib in dl.libs:
                    if lib not in dlopen_libnames:
                        libnames_full.append(f"{self.platform.libprefix}{lib}{libext}")
            if dl.dlopenlibs:
                libnames_full += [
                    f"{self.platform.libprefix}{lib}{libext}" for lib in dl.dlopenlibs
                ]

        return libnames_full

    def get_dlopen_library_names(self) -> Optional[List[str]]:
        libs = []
        if self.cfg.download:
            for dl in self.cfg.download:
                if dl.dlopenlibs:
                    libs += dl.dlopenlibs
        return libs

    def get_extra_objects(self) -> Optional[List[str]]:
        pass

    def get_type_casters_cfg(self, casters: Dict[str, Dict[str, Any]]) -> None:
        for ccfg in self.cfg.type_casters:
            cfg = {"hdr": ccfg.header}
            if ccfg.default_arg_cast:
                cfg["darg"] = True

            for typ in ccfg.types:
                casters[typ] = cfg

    def _update_addl_data_files(self) -> List[str]:
        headers = set()
        for ccfg in self.cfg.type_casters:
            headers.add(ccfg.header)

        if headers:
            includes = self.get_include_dirs()
            if includes:
                for hdr in headers:
                    for p in includes:
                        fpath = join(p, hdr)
                        if exists(fpath):
                            self._add_addl_data_file(fpath)

    def all_deps(self):
        if self._all_deps is None:
            self._all_deps = self.pkgcfg.get_all_deps(self.name)
        return self._all_deps

    def _all_includes(self, include_rpyb):
        includes = self.get_include_dirs()
        for dep in self.all_deps():
            dep_inc = dep.get_include_dirs()
            if dep_inc:
                includes.extend(dep_inc)
        if include_rpyb:
            includes.extend(self.pkgcfg.get_pkg("robotpy-build").get_include_dirs())
        return includes

    def _generation_search_path(self):
        return [self.root] + self._all_includes(False)

    def _all_library_dirs(self):
        libs = self.get_library_dirs()
        for dep in self.cfg.depends:
            libdirs = self.pkgcfg.get_pkg(dep).get_library_dirs()
            if libdirs:
                libs.extend(libdirs)
        return libs

    def _all_library_names(self):
        libs = list(
            set(self.get_library_names()) | set(self.get_dlopen_library_names())
        )
        for dep in self.cfg.depends:
            pkg = self.pkgcfg.get_pkg(dep)
            libnames = pkg.get_library_names()
            if libnames:
                libs.extend(libnames)
        return list(reversed(libs))

    def _all_extra_objects(self):
        libs = []
        for dep in self.cfg.depends:
            pkg = self.pkgcfg.get_pkg(dep)
            libnames = pkg.get_extra_objects()
            if libnames:
                libs.extend(libnames)
        return list(reversed(libs))

    def _all_casters(self):
        casters = {}
        for dep in self.all_deps():
            dep.get_type_casters_cfg(casters)
        self.pkgcfg.get_pkg("robotpy-build").get_type_casters_cfg(casters)
        self.get_type_casters_cfg(casters)

        # make each configuration unique
        for k, v in list(casters.items()):
            v = v.copy()
            v["typename"] = k
            casters[k] = v

        # add non-namespaced versions of all casters
        # -> in theory this could lead to a conflict, but
        #    let's see how it works in practice?
        for k, v in list(casters.items()):
            k = k.split("::")[-1]
            casters[k] = v
        return casters

    def on_build_dl(self, cache: str, srcdir: str):
        pkgcfgpy = join(self.root, "pkgcfg.py")
        srcdir = join(srcdir, self.name)

        try:
            os.unlink(self.libinit_import_py)
        except OSError:
            pass

        try:
            os.unlink(pkgcfgpy)
        except OSError:
            pass

        libnames_full = []
        all_libs = []
        downloads = self.cfg.download
        if downloads:
            libnames_full, all_libs = self._clean_and_download(downloads, cache, srcdir)

        self._write_libinit_py(libnames_full)
        self._write_pkgcfg_py(pkgcfgpy, libnames_full)

        return all_libs

    def _apply_patches(self, patches: List[PatchInfo], root: str):
        import patch

        for p in patches:
            patch_path = join(self.setup_root, normpath(p.patch))
            ps = patch.PatchSet()
            with open(patch_path, "rb") as fp:
                if not ps.parse(fp):
                    raise ValueError(f"Error parsing patch '{patch_path}'")

            if not ps.apply(strip=p.strip, root=root):
                raise ValueError(f"Error applying patch '{patch_path}' to '{root}'")

    def _clean_and_download(
        self, downloads: List[Download], cache: str, srcdir: str
    ) -> Tuple[List[str], List[str]]:
        libdir = join(self.root, "lib")
        incdir = join(self.root, "include")

        add_libdir = False
        add_incdir = False

        # Remove downloaded/generated artifacts first
        shutil.rmtree(libdir, ignore_errors=True)
        shutil.rmtree(incdir, ignore_errors=True)
        shutil.rmtree(srcdir, ignore_errors=True)

        dlopen_libnames = self.get_dlopen_library_names()
        libnames_full = []
        all_libs = []

        for dl in downloads:
            # extract the whole thing into a directory when using for sources
            if dl.sources is not None:
                download_and_extract_zip(dl.url, srcdir, cache)
                sources = [join(srcdir, normpath(s)) for s in dl.sources]
                self.extension.sources.extend(sources)
                if dl.patches:
                    self._apply_patches(dl.patches, srcdir)
            elif dl.sources is not None:
                raise ValueError("sources must be None if use_sources is False!")
            elif dl.patches is not None:
                raise ValueError("patches must be None if use_sources is False!")

            if dl.libs or dl.dlopenlibs:
                add_libdir = True
                extract_names = []
                os.makedirs(libdir)

                libext = dl.libexts.get(self.platform.libext, self.platform.libext)
                linkext = dl.linkexts.get(self.platform.linkext, self.platform.linkext)
                if dl.libs:
                    for lib in dl.libs:
                        if lib not in dlopen_libnames:
                            name = f"{self.platform.libprefix}{lib}{libext}"
                            libnames_full.append(name)
                            extract_names.append(name)
                            if libext != linkext:
                                extract_names.append(
                                    f"{self.platform.libprefix}{lib}{linkext}"
                                )

                if dl.dlopenlibs:
                    dlopen_libnames_full = [
                        f"{self.platform.libprefix}{lib}{libext}"
                        for lib in dl.dlopenlibs
                    ]
                    libnames_full += dlopen_libnames_full
                    extract_names += dlopen_libnames_full

                to = {
                    posixpath.join(dl.libdir, libname): join(libdir, libname)
                    for libname in extract_names
                }
                all_libs.extend(to.values())
            else:
                to = {}

            if dl.incdir is not None:
                to[dl.incdir] = self.incdir
                add_incdir = True

            download_and_extract_zip(dl.url, to, cache)

            if dl.header_patches:
                self._apply_patches(dl.header_patches, incdir)

        if add_incdir:
            for f in glob.glob(join(glob.escape(incdir), "**"), recursive=True):
                self._add_addl_data_file(f)

        if add_libdir:
            for f in glob.glob(join(glob.escape(libdir), "**"), recursive=True):
                self._add_addl_data_file(f)

        return libnames_full, all_libs

    def _write_libinit_py(self, libnames):
        # This file exists to ensure that any shared library dependencies
        # are loaded for the compiled extension

        init = inspect.cleandoc(
            """
        
        # This file is automatically generated, DO NOT EDIT
        # fmt: off

        from os.path import abspath, join, dirname, exists
        _root = abspath(dirname(__file__))

        ##IMPORTS##

        """
        )

        init += "\n"

        if libnames:
            if self.platform.os == "osx":
                init += "from ctypes import CDLL, RTLD_GLOBAL\n\n"
            else:
                init += "from ctypes import cdll\n\n"

            for libname in libnames:
                init += "try:\n"
                if self.platform.os == "osx":
                    init += f'    _lib = CDLL(join(_root, "lib", "{libname}"), mode=RTLD_GLOBAL)\n'
                else:
                    init += f'    _lib = cdll.LoadLibrary(join(_root, "lib", "{libname}"))\n'

                init += "except FileNotFoundError:\n"
                init += f'    if not exists(join(_root, "lib", "{libname}")):\n'
                init += f'        raise FileNotFoundError("{libname} was not found on your system. Is this package correctly installed?")\n'
                if self.platform.os == "windows":
                    init += f'    raise Exception("{libname} could not be loaded. Do you have Visual Studio C++ Redistributible 2019 installed?")\n\n'
                else:
                    init += f'    raise FileNotFoundError("{libname} could not be loaded. There is a missing dependency.")\n\n'
        imports = []
        for dep in self.cfg.depends:
            pkg = self.pkgcfg.get_pkg(dep)
            if pkg.libinit_import:
                imports.append(pkg.libinit_import)

        if imports:
            imports = "# runtime dependencies\nimport " + "\nimport ".join(imports)
        else:
            imports = ""

        init = init.replace("##IMPORTS##", imports)

        with open(self.libinit_import_py, "w") as fp:
            fp.write(init)

        self._add_addl_data_file(self.libinit_import_py)

    def _write_pkgcfg_py(self, fname, libnames_full):
        library_dirs = "[]"
        library_dirs_rel = []
        library_names = self.get_library_names()
        if library_names:
            library_dirs = '[join(_root, "lib")]'
            library_dirs_rel = ["lib"]

        deps = []
        for dep in self.cfg.depends:
            pkg = self.pkgcfg.get_pkg(dep)
            if not pkg.static_lib:
                deps.append(dep)

        # write pkgcfg.py
        pkgcfg = inspect.cleandoc(
            f"""
        # fmt: off
        # This file is automatically generated, DO NOT EDIT

        from os.path import abspath, join, dirname
        _root = abspath(dirname(__file__))

        libinit_import = "{self.libinit_import}"
        depends = {repr(deps)}
        pypi_package = {repr(self.pypi_package)}

        def get_include_dirs():
            return [join(_root, "include"), join(_root, "rpy-include")##EXTRAINCLUDES##]

        def get_library_dirs():
            return {library_dirs}

        def get_library_dirs_rel():
            return {repr(library_dirs_rel)}
        
        def get_library_names():
            return {repr(library_names)}

        def get_library_full_names():
            return {repr(libnames_full)}
        """
        )

        extraincludes = ""
        if self.cfg.extra_includes:
            # these are relative to the root of the project, need
            # to resolve the path relative to the package
            pth = join(*self.package_name.split("."))

            for h in self.cfg.extra_includes:
                h = '", "'.join(relpath(normpath(h), pth).split(sep))
                extraincludes += f', join(_root, "{h}")'

        pkgcfg = pkgcfg.replace("##EXTRAINCLUDES##", extraincludes)

        type_casters = {}
        self.get_type_casters_cfg(type_casters)
        if type_casters:
            pkgcfg += "\n\n"
            pkgcfg += inspect.cleandoc(
                f"""

            def get_type_casters_cfg(casters):
                casters.update({repr(type_casters)})

            def get_type_casters(casters):
                t = {{}}
                get_type_casters_cfg(t)
                for k, v in t.items():
                    if "hdr" in v:
                        casters[k] = v["hdr"]
            """
            )

        with open(fname, "w") as fp:
            fp.write(pkgcfg)

        self._add_addl_data_file(fname)

    def on_build_gen(
        self, cxx_gen_dir, missing_reporter: Optional[MissingReporter] = None
    ):
        if not self.cfg.autogen_headers:
            return

        cxx_gen_dir = join(cxx_gen_dir, self.name)

        if missing_reporter:
            report_only = True
        else:
            report_only = False
            missing_reporter = MissingReporter()

        hppoutdir = join(self.rpy_incdir, "rpygen")

        pp_includes = self._all_includes(True) + [sysconfig.get_path("include")]

        # TODO: only regenerate files if the generated files
        #       have changed
        if not report_only:
            if self.dev_config.only_generate is None:
                shutil.rmtree(cxx_gen_dir, ignore_errors=True)
                shutil.rmtree(hppoutdir, ignore_errors=True)

            os.makedirs(cxx_gen_dir, exist_ok=True)
            os.makedirs(hppoutdir, exist_ok=True)

        per_header = False
        data_fname = self.cfg.generation_data
        if self.cfg.generation_data:
            datapath = join(self.setup_root, normpath(self.cfg.generation_data))
            per_header = isdir(datapath)
            if not per_header:
                data = AutowrapConfigYaml.from_file(datapath)
        else:
            data = AutowrapConfigYaml()

        pp_defines = [self._cpp_version] + self.platform.defines + self.cfg.pp_defines
        casters = self._all_casters()

        # These are written to file to make it easier for dev mode to work
        classdeps = {}

        if self.dev_config.only_generate is not None:
            only_generate = {n: True for n in self.dev_config.only_generate}
        else:
            only_generate = None

        generation_search_path = self._generation_search_path()

        for name, header in self.cfg.autogen_headers.items():
            header = normpath(header)
            for path in generation_search_path:
                header_path = join(path, header)
                if exists(header_path):
                    header_root = pathlib.Path(path)
                    break
            else:
                import pprint

                pprint.pprint(generation_search_path)
                raise ValueError("could not find " + header)

            if not report_only:
                classdeps_dst = join(cxx_gen_dir, f"{name}.json")
                classdeps[name] = classdeps_dst

            if per_header:
                data_fname = join(datapath, name + ".yml")
                if not exists(data_fname):
                    print("WARNING: could not find", data_fname)
                    data = AutowrapConfigYaml()
                else:
                    data = AutowrapConfigYaml.from_file(data_fname)

            if only_generate is not None and not only_generate.pop(name, False):
                continue

            popts = ParserOptions(
                preprocessor=make_preprocessor(
                    defines=pp_defines,
                    include_paths=pp_includes,
                    encoding=data.encoding,
                )
            )

            gendata = GeneratorData(data)

            try:
                hctx = parse_header(
                    name,
                    pathlib.Path(header_path),
                    header_root,
                    gendata,
                    popts,
                    casters,
                    report_only,
                )

                if not report_only:
                    generated_sources = self.wwriter.write_files(
                        hctx, name, cxx_gen_dir, hppoutdir, classdeps_dst
                    )
                    self.extension.sources.extend(
                        [relpath(src, self.setup_root) for src in generated_sources]
                    )
            except Exception as e:
                raise ValueError(f"processing {header}") from e

            gendata.report_missing(data_fname, missing_reporter)

        if only_generate:
            unused = ", ".join(sorted(only_generate))
            # raise ValueError(f"only_generate specified unused headers! {unused}")
            # TODO: make this a warning

        if not report_only:
            for name, contents in missing_reporter.as_yaml():
                print("WARNING: some items not in generation yaml for", basename(name))
                print(contents)

        # generate an inline file that can be included + called
        if not report_only:
            self._write_wrapper_hpp(cxx_gen_dir, classdeps)
            gen_includes = [cxx_gen_dir]
        else:
            gen_includes = []

        self._gen_includes = gen_includes

        for f in glob.glob(join(glob.escape(hppoutdir), "*.hpp")):
            self._add_addl_data_file(f)

    def finalize_extension(self):
        if self.extension is None:
            return

        # Add the root to the includes (but only privately)
        root_includes = [self.root]

        # update the build extension so that build_ext works
        # use normpath to get rid of .. otherwise gcc is weird
        self.extension.include_dirs = [
            normpath(p)
            for p in (self._all_includes(True) + self._gen_includes + root_includes)
        ]
        self.extension.library_dirs = self._all_library_dirs()
        self.extension.libraries = self._all_library_names()
        self.extension.extra_objects = self._all_extra_objects()

    def _write_wrapper_hpp(self, outdir, classdeps):
        decls = []
        begin_calls = []
        finish_calls = []

        # Need to ensure that wrapper initialization is called in base order
        # so we have to toposort it here. The data is written at gen time
        # to JSON files
        types2name = {}
        types2deps = {}
        ordering = []

        for name, jsonfile in classdeps.items():
            with open(jsonfile) as fp:
                dep = json.load(fp)

            # make sure objects without classes are also included!
            if not dep:
                ordering.append(name)

            for clsname, bases in dep.items():
                if clsname in types2name:
                    raise ValueError(f"{name} ({jsonfile}): duplicate class {clsname}")
                types2name[clsname] = name
                types2deps[clsname] = bases[:]

        to_sort: Dict[str, Set[str]] = {}
        for clsname, bases in types2deps.items():
            clsname = types2name[clsname]
            deps = to_sort.setdefault(clsname, set())
            for base in bases:
                base = types2name.get(base)
                if base and base != clsname:
                    deps.add(base)

        ordering.extend(toposort.toposort_flatten(to_sort, sort=True))

        for name in ordering:
            decls.append(f"void begin_init_{name}(py::module &m);")
            decls.append(f"void finish_init_{name}();")
            begin_calls.append(f"    begin_init_{name}(m);")
            finish_calls.append(f"    finish_init_{name}();")

        content = (
            inspect.cleandoc(
                """

        // This file is autogenerated, DO NOT EDIT
        #pragma once
        #include <robotpy_build.h>

        // forward declarations
        ##DECLS##

        static void initWrapper(py::module &m) {
        ##BEGIN_CALLS##

        ##FINISH_CALLS##
        }
        
        """
            )
            .replace("##DECLS##", "\n".join(decls))
            .replace("##BEGIN_CALLS##", "\n".join(begin_calls))
            .replace("##FINISH_CALLS##", "\n".join(finish_calls))
        )

        with open(join(outdir, "rpygen_wrapper.hpp"), "w") as fp:
            fp.write(content)
