import fnmatch
import glob
import os
import shutil
from collections import defaultdict


class GlobDirectoryWalker:
    # a forward iterator that traverses a directory tree

    def __init__(self, directory, pattern="*"):
        self.stack = [directory]
        self.pattern = pattern
        self.files = []
        self.index = 0

    def __getitem__(self, index):
        while True:
            try:
                file = self.files[self.index]
                self.index = self.index + 1
            except IndexError:
                # pop next directory from stack
                self.directory = self.stack.pop()
                if os.path.isdir(self.directory):
                    self.files = os.listdir(self.directory)
                else:
                    self.files = []
                self.index = 0
            else:
                # got a filename
                fullname = os.path.join(self.directory, file)
                if os.path.isdir(fullname) and not os.path.islink(fullname):
                    self.stack.append(fullname)
                if fnmatch.fnmatch(file, self.pattern):
                    return fullname


def main(rootpath, destdir):
    if not os.path.exists(destdir):
        shutil.os.makedirs(destdir)

    examplesfnames = [fname for fname in GlobDirectoryWalker(rootpath, "*.py")]
    allowed_additional_file_extensions = [".rst", ".md", ".txt", ".mplstyle"]
    additional_files = [
        fname
        for fname in GlobDirectoryWalker(rootpath, "*.*")
        if os.path.splitext(fname)[1] in allowed_additional_file_extensions
    ]

    data_file_extensions = [".npy"]
    data_files = [
        fname
        for fname in GlobDirectoryWalker(rootpath, "*.*")
        if os.path.splitext(fname)[1] in data_file_extensions
    ]

    print(f"Documenting {len(examplesfnames)} examples")
    print(f"Documenting {len(additional_files)} additional files")
    print(f"Documenting {len(data_files)} data files")

    examplespaths = []
    examplesbasenames = []
    relativepaths = []
    outnames = []
    for f in examplesfnames:
        path, file = os.path.split(f)
        relpath = os.path.relpath(path, rootpath)
        if relpath == ".":
            relpath = ""
        path = os.path.normpath(path)
        filebase, ext = os.path.splitext(file)
        exname = filebase
        if relpath:
            exname = relpath.replace("/", ".").replace("\\", ".") + "." + exname
        examplespaths.append(path)
        examplesbasenames.append(filebase)
        relativepaths.append(relpath)
        outnames.append(exname)
    # We assume all files are encoded as UTF-8
    examplescode = []
    for fname in examplesfnames:
        with open(fname, encoding="utf-8") as f:
            examplescode.append(f.read())
    examplesdocs = []
    examplesafterdoccode = []
    for code in examplescode:
        codesplit = code.split("\n")
        comment_lines = 0
        for line in codesplit:
            if line.startswith("#") or len(line) == 0:
                comment_lines += 1
            else:
                break
        codesplit = codesplit[comment_lines:]
        readingdoc = False
        doc = []
        afterdoccode = ""
        for i in range(len(codesplit)):
            stripped = codesplit[i].strip()
            if stripped[:3] == '"""' or stripped[:3] == "'''":
                if not readingdoc:
                    readingdoc = True
                else:
                    afterdoccode = "\n".join(codesplit[i + 1 :])
                    break
            elif readingdoc:
                doc.append(codesplit[i])
            else:  # No doc
                afterdoccode = "\n".join(codesplit[i:])
                break
        examplesdocs.append("\n".join(doc))
        examplesafterdoccode.append(afterdoccode)

    categories = defaultdict(list)
    examples = zip(
        examplesfnames,
        examplespaths,
        examplesbasenames,
        examplescode,
        examplesdocs,
        examplesafterdoccode,
        relativepaths,
        outnames,
    )
    # Get the path relative to the examples director (not relative to the
    # directory where this file is installed
    if "BRIAN2_DOCS_EXAMPLE_DIR" in os.environ:
        rootdir = os.environ["BRIAN2_DOCS_EXAMPLE_DIR"]
    else:
        rootdir, _ = os.path.split(__file__)
        rootdir = os.path.normpath(os.path.join(rootdir, "../../examples"))
    eximgpath = os.path.abspath(
        os.path.join(rootdir, "../docs_sphinx/resources/examples_images")
    )
    print("Searching for example images in directory", eximgpath)
    for _fname, _path, basename, _code, docs, afterdoccode, relpath, exname in examples:
        categories[relpath].append((exname, basename))
        title = "Example: " + basename
        output = ".. currentmodule:: brian2\n\n"
        output += ".. " + basename + ":\n\n"
        output += title + "\n" + "=" * len(title) + "\n\n"
        note = f"""
        .. only:: html

            .. |launchbinder| image:: file:///usr/share/doc/python-brian-doc/docs/badge.svg
            .. _launchbinder: https://mybinder.org/v2/gh/brian-team/brian2-binder/master?filepath=examples/{exname.replace('.', '/')}.ipynb

            .. note::
               You can launch an interactive, editable version of this
               example without installing any local files
               using the Binder service (although note that at some times this
               may be slow or fail to open): |launchbinder|_

        """
        output += note + "\n\n"
        output += docs + "\n\n::\n\n"
        output += "\n".join(["    " + line for line in afterdoccode.split("\n")])
        output += "\n\n"

        eximgpattern = os.path.join(eximgpath, f"{exname}.*")
        images = glob.glob(eximgpattern + ".png") + glob.glob(eximgpattern + ".gif")
        for image in sorted(images):
            _, image = os.path.split(image)
            print("Found example image file", image)
            output += f".. image:: ../resources/examples_images/{image}\n\n"

        with open(os.path.join(destdir, exname + ".rst"), "w", encoding="utf-8") as f:
            f.write(output)

    category_additional_files = defaultdict(list)
    for fname in additional_files:
        path, file = os.path.split(fname)
        relpath = os.path.relpath(path, rootpath)
        if relpath == ".":
            relpath = ""
        full_name = relpath.replace("/", ".").replace("\\", ".") + "." + file + ".rst"
        category_additional_files[relpath].append((file, full_name))
        with open(fname, encoding="utf-8") as f:
            content = f.read()
        output = file + "\n" + "=" * len(file) + "\n\n"
        output += ".. code-block:: none\n\n"
        content_lines = ["\t" + line for line in content.split("\n")]
        output += "\n".join(content_lines)
        output += "\n\n"
        with open(os.path.join(destdir, full_name), "w", encoding="utf-8") as f:
            f.write(output)

    category_data_files = defaultdict(list)
    for fname in data_files:
        path, file = os.path.split(fname)
        relpath = os.path.relpath(path, rootpath)
        if relpath == ".":
            relpath = ""
        full_name = relpath.replace("/", ".").replace("\\", ".") + "." + file + ".rst"
        category_data_files[relpath].append((file, full_name))
        output = file + "\n" + "=" * len(file) + "\n\n"
        output += (
            "Download:"
            f" :download:`{file} <{os.path.join('/../examples/', relpath, file)}>`\n\n"
        )
        with open(os.path.join(destdir, full_name), "w", encoding="utf-8") as f:
            f.write(output)

    mainpage_text = "Examples\n"
    mainpage_text += "========\n\n"

    def insert_category(category, mainpage_text):
        if category:
            label = category.lower().replace(" ", "-").replace("/", ".")
            mainpage_text += f"\n.. _{label}:\n\n"
            mainpage_text += "\n" + category + "\n" + "-" * len(category) + "\n\n"
        mainpage_text += ".. toctree::\n"
        mainpage_text += "   :maxdepth: 1\n\n"
        for exname, basename in sorted(categories[category]):
            mainpage_text += f"   {basename} <{exname}>\n"
        for fname, full_name in sorted(category_additional_files[category]):
            mainpage_text += f"   {fname} <{full_name}>\n"
        for fname, full_name in sorted(category_data_files[category]):
            mainpage_text += f"   {fname} <{full_name}>\n"
        return mainpage_text

    mainpage_text = insert_category("", mainpage_text)
    for category in sorted(categories.keys()):
        if category:
            mainpage_text = insert_category(category, mainpage_text)

    with open(os.path.join(destdir, "index.rst"), "w") as f:
        f.write(mainpage_text)


if __name__ == "__main__":
    main("../../examples", "../../docs_sphinx/examples")