# Based on https://github.com/darwindarak/mdx_bib
import os
import re
from bisect import bisect_right
from typing import Tuple

import jedi
import mkdocs.structure.pages
import parso
import regex
from mkdocs.config.config_options import Type as MkType
from mkdocs.config.defaults import MkDocsConfig
from mkdocs.plugins import BasePlugin

from docs.scripts.autorefs.plugin import AutorefsPlugin

try:
    from importlib.metadata import entry_points
except ImportError:
    from importlib_metadata import entry_points


from bs4 import BeautifulSoup

BRACKET_RE = re.compile(r"\[([^\[]+)\]")
CITE_RE = re.compile(r"@([\w_:-]+)")
DEF_RE = re.compile(r"\A {0,3}\[@([\w_:-]+)\]:\s*(.*)")
INDENT_RE = re.compile(r"\A\t| {4}(.*)")

HREF_REGEX = (
    r"(?<=<\s*(?:a[^>]*href|img[^>]*src)=)"
    r'(?:"([^"]*)"|\'([^\']*)|[ ]*([^ =>]*)(?![a-z]+=))'
)
# Maybe find something less specific ?
PIPE_REGEX = r"(?<![a-zA-Z0-9._-])eds[.]([a-zA-Z0-9._-]*)(?![a-zA-Z0-9._-])"

HTML_PIPE_REGEX = r"""(?x)
(?<![a-zA-Z0-9._-])
<span[^>]*>eds<\/span>
<span[^>]*>[.]<\/span>
<span[^>]*>([a-zA-Z0-9._-]*)<\/span>
(?![a-zA-Z0-9._-])
"""

CITATION_RE = r"(\[@(?:[\w_:-]+)(?: *, *@(?:[\w_:-]+))*\])"


class ClickableSnippetsPlugin(BasePlugin):
    config_scheme: Tuple[Tuple[str, MkType]] = ()

    @mkdocs.plugins.event_priority(1000)
    def on_config(self, config: MkDocsConfig):
        for event_name, events in config.plugins.events.items():
            for event in list(events):
                if "autorefs" in str(event):
                    events.remove(event)
        old_plugin = config["plugins"]["autorefs"]
        plugin_config = dict(old_plugin.config)
        plugin = AutorefsPlugin()
        config.plugins["autorefs"] = plugin
        config["plugins"]["autorefs"] = plugin
        plugin.load_config(plugin_config)

    @classmethod
    def get_ep_namespace(cls, ep, namespace):
        if hasattr(ep, "select"):
            return ep.select(group=namespace)
        else:  # dict
            return ep.get(namespace, [])

    @mkdocs.plugins.event_priority(-1000)
    def on_post_page(
        self,
        output: str,
        page: mkdocs.structure.pages.Page,
        config: mkdocs.config.Config,
    ):
        """
        1. Replace absolute paths with path relative to the rendered page
           This must be performed after all other plugins have run.
        2. Replace component names with links to the component reference

        Parameters
        ----------
        output
        page
        config

        Returns
        -------

        """

        autorefs: AutorefsPlugin = config["plugins"]["autorefs"]
        ep = entry_points()
        page_url = os.path.join("/", page.file.url)
        spacy_factories_entry_points = {
            ep.name: ep.value
            for ep in (
                *self.get_ep_namespace(ep, "spacy_factories"),
                *self.get_ep_namespace(ep, "edsnlp_factories"),
                *self.get_ep_namespace(ep, "spacy_scorers"),
            )
        }

        def replace_component(match):
            full_group = match.group(0)
            name = "eds." + match.group(1)
            ep = spacy_factories_entry_points.get(name)
            preceding = output[match.start(0) - 50 : match.start(0)]
            if ep is not None and "DEFAULT:" not in preceding:
                try:
                    url = autorefs.get_item_url(ep.replace(":", "."))
                except KeyError:
                    pass
                else:
                    return f"<a href={url}>{name}</a>"
            return full_group

        def replace_link(match):
            relative_url = url = match.group(1) or match.group(2) or match.group(3)
            if url.startswith("/"):
                relative_url = os.path.relpath(url, page_url)
            return f'"{relative_url}"'

        output = regex.sub(PIPE_REGEX, replace_component, output)
        output = regex.sub(HTML_PIPE_REGEX, replace_component, output)

        all_snippets = ""
        all_offsets = []
        all_nodes = []

        soups = []

        # Replace absolute paths with path relative to the rendered page
        for match in regex.finditer("<code>.*?</code>", output, flags=regex.DOTALL):
            node = match.group(0)
            if "\n" in node:
                soup, snippet, python_offsets, html_nodes = self.convert_html_to_code(
                    node
                )
                size = len(all_snippets)
                all_snippets += snippet + "\n"
                all_offsets.extend([size + i for i in python_offsets])
                all_nodes.extend(html_nodes)
                soups.append((soup, match.start(0), match.end(0)))

        interpreter = jedi.Interpreter(all_snippets, [{}])
        line_lengths = [0]
        for line in all_snippets.split("\n"):
            line_lengths.append(len(line) + line_lengths[-1] + 1)
        line_lengths[-1] -= 1

        for name in self.iter_names(interpreter._module_node):
            try:
                line, col = name.start_pos
                offset = line_lengths[line - 1] + col
                node_idx = bisect_right(all_offsets, offset) - 1

                node = all_nodes[node_idx]
                gotos = interpreter.goto(line, col, follow_imports=True)
                gotos = [
                    goto
                    for goto in gotos
                    if (
                        goto
                        and goto.full_name
                        and goto.full_name.startswith("edsnlp")
                        and goto.type != "module"
                    )
                ]
                goto = gotos[0] if gotos else None
                if goto:
                    url = autorefs.get_item_url(goto.full_name)
                    # Check if node has no link in its upstream ancestors
                    if not node.find_parents("a"):
                        node.replace_with(
                            BeautifulSoup(
                                f'<a class="discrete-link" href="{url}">{node}</a>',
                                "html5lib",
                            )
                        )
            except Exception:
                pass

        # Re-insert soups into the output
        for soup, start, end in reversed(soups):
            output = output[:start] + str(soup.find("code")) + output[end:]

        output = regex.sub(HREF_REGEX, replace_link, output)

        return output

    @classmethod
    def iter_names(cls, root):
        if isinstance(root, parso.python.tree.Name):
            yield root
        for child in getattr(root, "children", ()):
            yield from cls.iter_names(child)

    @classmethod
    def convert_html_to_code(
        cls, html_content: str
    ) -> Tuple[BeautifulSoup, str, list, list]:
        pre_html_content = "<pre>" + html_content + "</pre>"
        soup = list(BeautifulSoup(pre_html_content, "html5lib").children)[0]
        code_element = soup.find("code")

        line_lengths = [0]
        for line in pre_html_content.split("\n"):
            line_lengths.append(len(line) + line_lengths[-1] + 1)
        line_lengths[-1] -= 1

        python_code = ""
        code_offsets = []
        html_nodes = []
        code_offset = 0

        def extract_text_with_offsets(el):
            nonlocal python_code, code_offset
            for content in el.contents:
                # check not class md-annotation
                # Recursively process child elements
                if isinstance(content, str):
                    python_code += content
                    code_offsets.append(code_offset)
                    code_offset += len(content)
                    html_nodes.append(content)
                    continue
                if "md-annotation" not in content.get("class", ""):
                    extract_text_with_offsets(content)

        extract_text_with_offsets(code_element)

        return soup, python_code, code_offsets, html_nodes
