import ast

import inspect

import re

import sys

import textwrap

import warnings

import catalogue

import pytest

from spacy.tokens.underscore import Underscore

pytest.importorskip("mkdocs")

try:

    import torch.nn

except ImportError:

    torch = None

if torch is None:

    pytest.skip("torch not installed", allow_module_level=True)

pytest.importorskip("rich")

from extract_docs_code import extract_docs_code  # noqa: E402

# We don't check documentation for Python <= 3.7:

if sys.version_info < (3, 8):

    url_to_code = {}

else:

    url_to_code = dict(extract_docs_code())

    # just to make sure something didn't go wrong

    assert len(url_to_code) > 50

def printer(code: str) -> None:

    """

    Prints a code bloc with lines for easier debugging.

    Parameters

    ----------

    code : str

        Code bloc.

    """

    lines = []

    for i, line in enumerate(code.split("\n")):

        lines.append(f"{i + 1:03}  {line}")

    print("\n".join(lines))

def insert_assert_statements(code):

    line_table = [0]

    for line in code.splitlines(keepends=True):

        line_table.append(line_table[-1] + len(line))

    tree = ast.parse(code)

    replacements = []

    for match in re.finditer(

        r"^\s*#\s*Out\s*: (.*$(?:\n#\s.*$)*)", code, flags=re.MULTILINE

    ):

        lineno = code[: match.start()].count("\n")

        for stmt in tree.body:

            if stmt.end_lineno == lineno:

                if isinstance(stmt, ast.Expr):

                    expected = textwrap.dedent(match.group(1)).replace("\n# ", "\n")

                    begin = line_table[stmt.lineno - 1]

                    if not (expected.startswith("'") or expected.startswith('"')):

                        expected = repr(expected)

                    end = match.end()

                    stmt_str = ast.unparse(stmt)

                    if stmt_str.startswith("print("):

                        stmt_str = stmt_str[len("print") :]

                    repl = f"""\

value = {stmt_str}

assert {expected} == str(value)

"""

                    replacements.append((begin, end, repl))

                if isinstance(stmt, ast.For):

                    expected = textwrap.dedent(match.group(1)).split("\n# Out: ")

                    expected = [line.replace("\n# ", "\n") for line in expected]

                    begin = line_table[stmt.lineno - 1]

                    end = match.end()

                    stmt_str = ast.unparse(stmt).replace("print", "assert_print")

                    repl = f"""\

printed = []

{stmt_str}

assert {expected} == printed

"""

                    replacements.append((begin, end, repl))

    for begin, end, repl in reversed(replacements):

        code = code[:begin] + repl + code[end:]

    return code

# TODO: once in a while, it can be interesting to run reset_imports for each code block,

#  instead of only once and tests should still pass, but it's way slower.

@pytest.fixture(scope="module")

def reset_imports():

    """

    Reset the imports for each test.

    """

    # 1. Clear registered functions to avoid using cached ones

    for k, m in list(catalogue.REGISTRY.items()):

        mod = inspect.getmodule(m)

        if mod is not None and mod.__name__.startswith("edsnlp"):

            del catalogue.REGISTRY[k]

    # Let's ensure that we "bump" into every possible warnings:

    # 2. Remove all modules that start with edsnlp, to reimport them

    for k in list(sys.modules):

        if k.split(".")[0] == "edsnlp":

            del sys.modules[k]

    # 3. Delete spacy extensions to avoid error when re-importing

    Underscore.span_extensions.clear()

    Underscore.doc_extensions.clear()

    Underscore.token_extensions.clear()

# Note the use of `str`, makes for pretty output

@pytest.mark.parametrize("url", sorted(url_to_code.keys()), ids=str)

def test_code_blocks(url, tmpdir, reset_imports):

    code = url_to_code[url]

    code_with_asserts = """

def assert_print(*args, sep=" ", end="\\n", file=None, flush=False):

    printed.append((sep.join(map(str, args)) + end).rstrip('\\n'))

""" + insert_assert_statements(code)

    assert "# Out:" not in code_with_asserts, (

        "Unparsed asserts in {url}:\n" + code_with_asserts

    )

    # We'll import test_code_blocks from here

    sys.path.insert(0, str(tmpdir))

    test_file = tmpdir.join("test_code_blocks.py")

    # Clear all warnings

    warnings.resetwarnings()

    try:

        with warnings.catch_warnings():

            warnings.simplefilter("error")

            warnings.filterwarnings(module=".*endlines.*", action="ignore")

            warnings.filterwarnings(

                message="__package__ != __spec__.parent", action="ignore"

            )

            # First, forget test_code_blocks

            sys.modules.pop("test_code_blocks", None)

            # Then, reimport it, to let pytest do its assertion rewriting magic

            test_file.write_text(code_with_asserts, encoding="utf-8")

            import test_code_blocks  # noqa: F401

            exec(

                compile(code_with_asserts, test_file, "exec"),

                {"__MODULE__": "__main__"},

            )

    except Exception:

        printer(code_with_asserts)

        raise