import re

import shutil

import tempfile

from textwrap import dedent

from typing import Tuple

from markdown.extensions import Extension

from markdown.extensions.attr_list import get_attrs

from markdown.extensions.codehilite import parse_hl_lines

from markdown.extensions.fenced_code import FencedBlockPreprocessor

from mkdocs.commands.build import build

from mkdocs.config import load_config

from mkdocs.config.config_options import Type as MkType

from mkdocs.config.defaults import MkDocsConfig

from mkdocs.plugins import BasePlugin

from mkdocstrings.extension import AutoDocProcessor

from mkdocstrings.plugin import MkdocstringsPlugin

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}(.*)")

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

class PyCodePreprocessor(FencedBlockPreprocessor):

    """Gather reference definitions and citation keys"""

    FENCED_BLOCK_RE = re.compile(

        dedent(

            r"""

            (?P<fence>^[ ]*(?:~{3,}|`{3,}))[ ]*                          # opening fence

            ((\{(?P<attrs>[^\}\n]*)\})|                              # (optional {attrs} or

            (\.?(?P<lang>[\w#.+-]*)[ ]*)?                            # optional (.)lang

            (hl_lines=(?P<quot>"|')(?P<hl_lines>.*?)(?P=quot)[ ]*)?) # optional hl_lines)

            \n                                                       # newline (end of opening fence)

            (?P<code>.*?)(?<=\n)                                     # the code block

            (?P=fence)[ ]*$                                          # closing fence

        """  # noqa: E501

        ),

        re.MULTILINE | re.DOTALL | re.VERBOSE,

    )

    def __init__(self, md, code_blocks):

        super().__init__(md, {})

        self.code_blocks = code_blocks

    def run(self, lines):

        text = "\n".join(lines)

        while True:

            # ----  https://github.com/Python-Markdown/markdown/blob/5a2fee/markdown/extensions/fenced_code.py#L84C9-L98  # noqa: E501

            m = self.FENCED_BLOCK_RE.search(text)

            if m:

                lang, id, classes, config = None, "", [], {}

                if m.group("attrs"):

                    id, classes, config = self.handle_attrs(get_attrs(m.group("attrs")))

                    if len(classes):

                        lang = classes.pop(0)

                else:

                    if m.group("lang"):

                        lang = m.group("lang")

                    if m.group("hl_lines"):

                        # Support `hl_lines` outside of `attrs` for

                        # backward-compatibility

                        config["hl_lines"] = parse_hl_lines(m.group("hl_lines"))

                # ----

                code = m.group("code")

                if lang == "python" and "no-check" not in classes:

                    self.code_blocks.append(dedent(code))

            else:

                break

            text = text[m.end() :]

        return lines

context_citations = None

class PyCodeExtension(Extension):

    def __init__(self, code_blocks):

        super(PyCodeExtension, self).__init__()

        self.code_blocks = code_blocks

    def extendMarkdown(self, md):

        self.md = md

        md.registerExtension(self)

        md.preprocessors.register(

            PyCodePreprocessor(md, self.code_blocks), "fenced_code", 31

        )

        for ext in md.registeredExtensions:

            if isinstance(ext, AutoDocProcessor):

                ext._config["mdx"].append(self)

def makeExtension(*args, **kwargs):

    return PyCodeExtension(*args, **kwargs)

class PyCodeExtractorPlugin(BasePlugin):

    config_scheme: Tuple[Tuple[str, MkType]] = (

        # ("bibtex_file", MkType(str)),  # type: ignore[assignment]

        # ("order", MkType(str, default="unsorted")),  # type: ignore[assignment]

    )

    def __init__(self, global_config):

        self.global_config = global_config

        self.page_code_blocks = []

        self.docs_code_blocks = []

    def on_config(self, config: MkDocsConfig):

        self.ext = PyCodeExtension(self.page_code_blocks)

        # After pymdownx.highlight, because of weird registering deleting the first

        # extension

        config["markdown_extensions"].append(self.ext)

        config["markdown_extensions"].remove("pymdownx.highlight")

        config["markdown_extensions"].remove("fenced_code")

    def on_pre_build(self, *, config: MkDocsConfig):

        mkdocstrings_plugin: MkdocstringsPlugin = config.plugins["mkdocstrings"]

        mkdocstrings_plugin.get_handler("python")

    def on_page_content(self, html, page, config, files):

        if len(self.page_code_blocks):

            self.docs_code_blocks.append((page.url, "\n".join(self.page_code_blocks)))

        self.page_code_blocks.clear()

        return html

def extract_docs_code():

    config = load_config()

    temp_dir = tempfile.mkdtemp()

    try:

        config["site_dir"] = temp_dir

        # plug the pycode extractor plugin

        plugin = PyCodeExtractorPlugin(config)

        config.plugins["pycode_extractor"] = plugin

        config["plugins"].run_event("startup", command="build", dirty=False)

        try:

            build(config)

        finally:

            config["plugins"].run_event("shutdown")

    finally:

        shutil.rmtree(temp_dir, ignore_errors=True)

    docs_code_blocks = plugin.docs_code_blocks

    # Deduplicate both keys and values

    return {

        k: v

        for v, k in {v: k for k, v in docs_code_blocks if "changelog" not in k}.items()

    }