# Copied from the yt_project, commit e8fb57e
# yt/doc/extensions/notebook_sphinxext.py
#  https://bitbucket.org/yt_analysis/yt/src/e8fb57e66ca42e26052dadf054a5c782740abec9/doc/extensions/notebook_sphinxext.py?at=yt

# Almost completely re-written by Matthew Harrigan to use nbconvert v4


import os
import shutil

import nbformat
from docutils import nodes
from docutils.parsers.rst import Directive, directives
from nbconvert import HTMLExporter, PythonExporter


def _read(wd, name):
    with open(f"{wd}/{name}.ipynb") as f:
        notebook = nbformat.read(f, as_version=4)
    return notebook


def export_html(wd, name):
    nb = _read(wd, name)

    config = {
        "Exporter": {
            "template_file": "embed",
            "template_path": ["./sphinxext/"],
        },
        "ExecutePreprocessor": {"enabled": True},
        "ExtractOutputPreprocessor": {"enabled": True},
        "CSSHTMLHeaderPreprocessor": {"enabled": True},
    }

    exporter = HTMLExporter(config)

    try:
        body, resources = exporter.from_notebook_node(nb)

        for fn, data in resources["outputs"].items():
            with open(f"{wd}/{fn}", "wb") as f:
                f.write(data)
        return body
    except Exception as e:
        return str(e)


def export_python(wd, name):
    nb = _read(wd, name)
    exporter = PythonExporter()
    body, resources = exporter.from_notebook_node(nb)
    with open(f"{wd}/{name}.py", "w") as f:
        f.write(body)


class NotebookDirective(Directive):
    """Insert an evaluated notebook into a document"""

    required_arguments = 1
    optional_arguments = 1
    option_spec = {"skip_exceptions": directives.flag}
    final_argument_whitespace = True

    def run(self):
        # check if raw html is supported
        if not self.state.document.settings.raw_enabled:
            raise self.warning('"%s" directive disabled.' % self.name)

        # get path to notebook
        nb_rel_path = self.arguments[0]
        nb_abs_path = f"{setup.confdir}/../{nb_rel_path}"
        nb_abs_path = os.path.abspath(nb_abs_path)
        nb_name = os.path.basename(nb_rel_path).split(".")[0]
        dest_dir = f"{setup.app.builder.outdir}/{os.path.dirname(nb_rel_path)}/{nb_name}"
        fmt = {"wd": dest_dir, "name": nb_name}

        if not os.path.exists(dest_dir):
            os.makedirs(dest_dir)

        shutil.copyfile(nb_abs_path, "{wd}/{name}.ipynb".format(**fmt))

        # TODO: Actually save evaluated notebook
        shutil.copyfile(nb_abs_path, "{wd}/{name}_eval.ipynb".format(**fmt))

        html = export_html(**fmt)
        export_python(**fmt)

        # Create link to notebook and script files
        link_rst = "({uneval}; {eval}; {py})".format(
            uneval=formatted_link("{wd}/{name}.ipynb".format(**fmt)),
            eval=formatted_link("{wd}/{name}_eval.ipynb".format(**fmt)),
            py=formatted_link("{wd}/{name}.py".format(**fmt)),
        )

        rst_file = self.state_machine.document.attributes["source"]
        self.state_machine.insert_input([link_rst], rst_file)

        # create notebook node
        attributes = {"format": "html", "source": "nb_path"}
        nb_node = notebook_node("", html, **attributes)
        nb_node.source, nb_node.line = self.state_machine.get_source_and_line(
            self.lineno,
        )

        # add dependency
        self.state.document.settings.record_dependencies.add(nb_abs_path)

        return [nb_node]


class notebook_node(nodes.raw):
    pass


def formatted_link(path):
    return f"`{os.path.basename(path)} <{path}>`__"


def visit_notebook_node(self, node):
    self.visit_raw(node)


def depart_notebook_node(self, node):
    self.depart_raw(node)


def setup(app):
    setup.app = app
    setup.config = app.config
    setup.confdir = app.confdir

    app.add_node(
        notebook_node,
        html=(visit_notebook_node, depart_notebook_node),
    )

    app.add_directive("notebook", NotebookDirective)