from __future__ import annotations

import importlib.util
import os
import platform
import random
import re
import subprocess as sp
import sys
import traceback
import uuid
from argparse import ArgumentParser
from contextlib import ExitStack
from contextlib import contextmanager
from io import StringIO
from pathlib import Path
from tempfile import TemporaryDirectory
from typing import Callable
from typing import ContextManager
from typing import Sequence
from unittest.mock import patch

from rich.console import Console
from rich.panel import Panel
from rich.text import Text

from inline_snapshot._exceptions import UsageError
from inline_snapshot._get_snapshot_value import get_snapshot_value
from inline_snapshot._snapshot_arg import snapshot_arg
from inline_snapshot._snapshot_session import SnapshotSession

from .._global_state import enter_snapshot_context
from .._global_state import leave_snapshot_context
from .._global_state import state
from .._types import Category
from .._types import Snapshot
from .._types import SnapshotArg
from ..extra import _format_exception

ansi_escape = re.compile(
    r"""
    \x1B  # ESC
    (?:   # 7-bit C1 Fe (except CSI)
        [@-Z\\-_]
    |     # or [ for CSI, followed by a control sequence
        \[
        [0-?]*  # Parameter bytes
        [ -/]*  # Intermediate bytes
        [@-~]   # Final byte
    )
""",
    re.VERBOSE,
)

ansi_osc_escape = re.compile(r"\x1B\].*?(?:\x07|\x1B\\)")


def normalize(text):
    text = ansi_osc_escape.sub("", text)
    text = ansi_escape.sub("", text)

    # fix windows problems
    text = text.replace("\u2500", "-")
    text = text.replace("\u2502", "|")
    for c in "┌┐└┘\u256d\u256e\u2570\u256f":
        text = text.replace(c, "+")
    text = text.replace("\r", "")
    return text


@contextmanager
def deterministic_uuid(seed=0):
    rd = random.Random(seed)

    def f():
        return uuid.UUID(int=rd.getrandbits(128), version=4)

    with patch("uuid.uuid4", new=f):
        yield


# this code is copied from pytest

# Regex to match the session duration string in the summary: "74.34s".
rex_session_duration = re.compile(r"\d+\.\d\ds")
# Regex to match all the counts and phrases in the summary line: "34 passed, 111 skipped".
rex_outcome = re.compile(r"(\d+) (\w+)")


def parse_outcomes(lines):
    for line in reversed(lines):
        if rex_session_duration.search(line):
            outcomes = rex_outcome.findall(line)
            ret = {noun: int(count) for (count, noun) in outcomes}
            break
        else:
            pass  # pragma: no cover
    else:
        return {}

    to_plural = {
        "warning": "warnings",
        "error": "errors",
    }
    return {to_plural.get(k, k): v for k, v in ret.items()}


def _pytest_error_line(line: str) -> str | None:
    if line and line.lstrip()[:2] in ("> ", "E "):
        return line
    return None


@contextmanager
def chdir(path):
    cwd = os.getcwd()
    try:
        os.chdir(path)
        yield
    finally:
        os.chdir(cwd)


console = Console(width=80)


def _subprocess_env() -> dict[str, str]:
    env_keys = [
        "PATH",
        "PWD",
        "VIRTUAL_ENV",
        "TOP",
        "COVERAGE_PROCESS_START",
        "PYTHONDONTWRITEBYTECODE",
        "PYTHONIOENCODING",
        "PYTHONPATH",
        "PYTHONWARNINGS",
        "PYTEST_DISABLE_PLUGIN_AUTOLOAD",
        "PYTEST_PLUGINS",
    ]
    if platform.system() == "Windows":  # pragma: no cover
        env_keys.extend(
            [
                "APPDATA",
                "COMSPEC",
                "LOCALAPPDATA",
                "PATHEXT",
                "PROGRAMDATA",
                "PROGRAMFILES",
                "PROGRAMFILES(X86)",
                "SYSTEMDRIVE",
                "SYSTEMROOT",
                "TEMP",
                "TMP",
                "USERPROFILE",
                "WINDIR",
            ]
        )

    return {key: os.environ[key] for key in env_keys if key in os.environ}


@contextmanager
def change_file(path: Path, map_function):
    exists = path.exists()
    if exists:
        text = path.read_bytes().decode("utf-8")
    else:
        text = ""

    path.write_bytes(map_function(text).encode("utf-8"))

    yield

    if exists:
        path.write_bytes(text.encode("utf-8"))
    else:
        path.unlink()


@contextmanager
def temp_environ(**kwargs):
    original = dict(os.environ)
    os.environ.update(kwargs)
    try:
        yield
    finally:
        os.environ.clear()
        os.environ.update(original)


class StopTesting(Exception):
    pass


class TestingException(Exception):
    def __init__(self, e):
        self.exception = e


class Example:
    files: dict[str, str | bytes]
    seed: int

    def __init__(self, files: str | bytes | dict[str, str | bytes]):
        """
        Parameters:
            files: a collection of files which are used as your example project,
                   or just a string which will be saved as *tests/test_something.py*.
        """
        if isinstance(files, (str, bytes)):
            files = {"tests/test_something.py": files}

        self.files = files
        self.seed = 0

    def _new_example(self, files: dict[str, str | bytes]) -> Example:
        e = Example(files)
        e.seed = self.seed
        return e

    def _next_seed(self):
        self.seed += 1
        return self

    def dump_files(self):

        for name, content in self.files.items():
            if isinstance(content, bytes):
                content = repr(content)
            console.print(Panel(Text(content), title=name))

    def _write_files(self, dir: Path):
        for name, content in self.files.items():
            filename = dir / name
            filename.parent.mkdir(exist_ok=True, parents=True)
            if isinstance(content, str):
                content = content.encode("utf-8")
            filename.write_bytes(content)

    def conftest_footer(self):
        return f"""
import uuid
import random

rd = random.Random({self.seed})

def f():
    return uuid.UUID(int=rd.getrandbits(128), version=4)

uuid.uuid4 = f
"""

    def _read_files(self, dir: Path):

        def try_read(path: Path) -> str | bytes:
            code = path.read_bytes()
            try:
                return code.decode("utf-8")
            except UnicodeDecodeError:
                return code

        def normalize_path(path):
            return str(path.relative_to(dir)).replace("\\", "/")

        return {
            normalize_path(p): try_read(p)
            for p in dir.rglob("*")
            if p.is_file()
            and p.name != ".gitignore"
            and p.suffix != ".pyc"
            and ".pytest_cache" not in p.parts
        }

    def with_files(self, extra_files: dict[str, str | bytes]) -> Example:
        """
        Adds extra files to the example.

        Arguments:
            extra_files: dictionary of filenames and file contents.
        """
        return self._new_example({**self.files, **extra_files})

    def read_file(self, filename: str) -> str:
        """
        Reads a file from the example.

        Arguments:
            filename: the filename.
        """
        text = self.files[filename]
        assert isinstance(text, str)
        return text

    read_text = read_file

    def change_code(
        self, mapping: Callable[[str], str], *, filename: str | None = None
    ) -> Example:
        """
        Changes example tests by mapping every file (if filename is None) with the given function.

        Arguments:
            mapping: function to apply to each file's content.
            filename: if given, only this file is changed
        """
        return self._new_example(
            {
                name: (
                    mapping(text)
                    if isinstance(text, str) and (name == filename or filename is None)
                    else text
                )
                for name, text in self.files.items()
            }
        )

    def replace(self, old_text: str, new_text: str) -> Example:
        """
        Changes example tests by replacing old_text with new_text.

        Arguments:
            old_text: the text to be replaced.
            new_text: the new text to use instead.
        """
        return self.change_code(lambda code: code.replace(old_text, new_text))

    def remove_file(self, filename: str) -> Example:
        """
        Removes a file from the example.

        Arguments:
            filename: the file to be removed.
        """
        return self._new_example(
            {name: file for name, file in self.files.items() if name != filename}
        )

    def move_file(self, old_filename: str, new_filename: str) -> Example:
        """
        Moves a file in the example.

        Arguments:
            old_filename: the file to be moved.
            new_filename: the new location of the file.
        """
        return self._new_example(
            {
                (new_filename if name == old_filename else name): file
                for name, file in self.files.items()
            }
        )

    def format(self, filename: str = "tests/test_something.py") -> Example:
        """
        Formats a Python file using black.

        Arguments:
            filename: the file to format (default: "tests/test_something.py").
        """
        from inline_snapshot._format import format_code

        content = self.files[filename]
        assert isinstance(content, str)

        # Write files to temp directory so black can find pyproject.toml
        with TemporaryDirectory() as dir:
            tmp_path = Path(dir)
            self._write_files(tmp_path)
            file_path = tmp_path / filename
            formatted = format_code(content, file_path)

        return self._new_example({**self.files, filename: formatted})

    def is_formatted(self, filename: str = "tests/test_something.py") -> bool:
        """
        Checks if a Python file is properly formatted with black.

        Arguments:
            filename: the file to check (default: "tests/test_something.py").
        """
        from inline_snapshot._format import format_code

        content = self.files[filename]
        assert isinstance(content, str)

        # Write files to temp directory so black can find pyproject.toml
        with TemporaryDirectory() as dir:
            tmp_path = Path(dir)
            self._write_files(tmp_path)
            file_path = tmp_path / filename
            formatted = format_code(content, file_path)

        return content == formatted

    def run_inline(
        self,
        args: list[str] = [],
        *,
        context_managers: Sequence[ContextManager] = (),
        reported_categories: SnapshotArg[set[Category]] | None = None,
        changed_files: SnapshotArg[dict[str, str]] = {},
        report: Snapshot[str] | None = None,
        raises: SnapshotArg[str] = "<no exception>",
        stderr: SnapshotArg[str] = "",
    ) -> Example:
        """Execute the example files in process and run every `test_*`
        function.

        This is useful for fast test execution.

        Parameters:
            args: inline-snapshot arguments (supports only "--inline-snapshot=fix|create|...").
            context_managers: list of context managers to use when the code is executed.
            reported_categories: snapshot of categories which inline-snapshot thinks could be applied or `None` when they are equal to the CLI flags.
            changed_files: snapshot of files which are changed by this run.
            raises: snapshot of the exception raised during test execution.
                    Required if your code raises an exception.

        Returns:
            A new Example instance containing the changed files.
        """

        parser = ArgumentParser()

        parser.add_argument(
            "--inline-snapshot",
            metavar="(disable,short-report,report,review,create,update,trim,fix)*",
            dest="inline_snapshot",
            help="update specific snapshot values:\n"
            "disable: disable the snapshot logic\n"
            "short-report: show a short report\n"
            "report: show a longer report with a diff report\n"
            "review: allow to approve the changes interactive\n"
            "create: creates snapshots which are currently not defined\n"
            "update: update snapshots even if they are defined\n"
            "trim: changes the snapshot in a way which will make the snapshot more precise.\n"
            "fix: change snapshots which currently break your tests\n",
        )
        parsed_args = parser.parse_args(args)
        flags = (parsed_args.inline_snapshot or "").split(",")
        self.dump_files()
        flags = {f for f in flags if f}

        with TemporaryDirectory() as dir:
            tmp_path = Path(dir)

            self._write_files(tmp_path)

            raised_exception = []

            with ExitStack() as stack:
                stack.enter_context(deterministic_uuid(self.seed))
                stack.enter_context(chdir(tmp_path))
                stack.enter_context(temp_environ(TERM="unknown"))

                session = SnapshotSession()

                def report_error(message):
                    raise StopTesting(message)

                snapshot_flags = set()
                old_modules = sys.modules
                old_path = sys.path[:]
                # Add tmp_path to sys.path so modules can be imported normally
                sys.path.insert(0, str(tmp_path))

                report_output = StringIO()
                console = Console(file=report_output, width=80)

                stderr_output = ""

                try:
                    for cm in context_managers:
                        stack.enter_context(cm)

                    enter_snapshot_context()
                    session.load_config(
                        tmp_path / "pyproject.toml",
                        flags,
                        parallel_run=False,
                        error=report_error,
                        project_root=tmp_path,
                    )

                    # Load and register all conftest.py files first
                    for conftest_path in tmp_path.rglob("conftest.py"):
                        print("load> conftest", conftest_path)

                        # Load conftest module using importlib
                        spec = importlib.util.spec_from_file_location(
                            f"conftest_{conftest_path.parent.name}", conftest_path
                        )
                        assert spec and spec.loader

                        conftest_module = importlib.util.module_from_spec(spec)
                        sys.modules[spec.name] = conftest_module
                        conftest_module.__file__ = str(conftest_path)
                        try:
                            spec.loader.exec_module(conftest_module)
                        except Exception as e:
                            raise TestingException(e) from e

                        # Register customize hooks from this conftest
                        session.register_customize_hooks_from_module(conftest_module)

                    tests_found = False
                    for filename in tmp_path.rglob("test_*.py"):
                        print("run> pytest-inline", filename, *args)

                        # Load module using importlib
                        spec = importlib.util.spec_from_file_location(
                            filename.stem, filename
                        )
                        assert spec and spec.loader

                        module = importlib.util.module_from_spec(spec)
                        sys.modules[filename.stem] = module
                        try:
                            spec.loader.exec_module(module)
                        except Exception as e:
                            raise TestingException(e) from e

                        # run all test_* functions
                        tests = [
                            v
                            for k, v in module.__dict__.items()
                            if (k.startswith("test_") or k == "test") and callable(v)
                        ]
                        tests_found |= len(tests) != 0

                        for v in tests:
                            try:

                                def fail(message):
                                    console.print(f"FAIL: {message}")

                                session.test_enter()
                                try:
                                    v()
                                finally:
                                    session.test_exit(fail=fail)
                            except Exception as e:
                                raise TestingException(e) from e

                    if not tests_found:
                        raise UsageError("no test_*() functions in the example")

                except StopTesting as e:
                    stderr_output = f"ERROR: {e}\n"
                except TestingException as e:
                    traceback.print_exc()
                    raised_exception.append(e.exception)
                finally:
                    try:
                        session.show_report(console)

                        for snapshot in state().snapshots.values():
                            for change in snapshot._changes():
                                snapshot_flags.add(change.flag)
                    except Exception as e:
                        traceback.print_exc()
                        raised_exception.append(e)

                    sys.modules = old_modules
                    sys.path = old_path
                    leave_snapshot_context()

            # make the assertions in the original context

            assert snapshot_arg(stderr) == stderr_output

            reported_categories = snapshot_arg(reported_categories)

            if snapshot_flags != flags:
                assert set(snapshot_flags) == reported_categories
            else:
                if get_snapshot_value(reported_categories) != set(snapshot_flags):
                    assert None == reported_categories

            raises = snapshot_arg(raises)

            if raised_exception:
                raises_text = "\n".join(_format_exception(e) for e in raised_exception)
                raises_text = raises_text.replace(str(tmp_path), "<tmp>")
                assert raises == raises_text
            else:
                assert raises == "<no exception>"

            assert snapshot_arg(changed_files) == self._changed_files(tmp_path)

            if report is not None:
                assert report == normalize(report_output.getvalue())

            return self._new_example(self._read_files(tmp_path))._next_seed()

    def _changed_files(self, tmp_path):
        current_files = {}
        all_current_files = self._read_files(tmp_path)

        for name, content in sorted(all_current_files.items()):
            if name not in self.files or self.files[name] != content:
                current_files[name] = content

        for name in self.files:
            if name not in all_current_files:
                current_files[name] = None

        return current_files

    def run_pytest(
        self,
        args: list[str] = [],
        *,
        term_columns=80,
        env: dict[str, str] = {},
        changed_files: SnapshotArg[dict[str, str]] = {},
        report: Snapshot[str] | None = None,
        error: SnapshotArg[str] = "",
        stderr: SnapshotArg[str] = "",
        stdout: SnapshotArg[str] | None = None,
        returncode: SnapshotArg[int] = 0,
        stdin: bytes = b"",
        outcomes: SnapshotArg[dict[str, int]] = {"passed": 1},
    ) -> Example:
        """Run pytest with the given args and environment variables in a separate
        process.

        It can be used to test the interaction between your code and pytest, but it is a bit slower than `run_inline`.

        Parameters:
            args: pytest arguments like "--inline-snapshot=fix"
            env: dict of environment variables
            changed_files: snapshot of files changed by this run.
            report: snapshot of the report at the end of the pytest run.
            stderr: pytest stderr output
            stdout: raw pytest stdout output with ansi escape sequences (this is still **experimental** and the datatype might change or it might be removed again in the future).
            returncode: snapshot of the pytest return code.

        Returns:
            A new Example instance containing the changed files.
        """

        self.dump_files()

        with TemporaryDirectory() as dir:
            tmp_path = Path(dir)

            self._write_files(tmp_path)

            pytest_args = list(args)
            pytest_args = ["--color=yes", *pytest_args]

            cmd = [sys.executable, "-m", "pytest", "-p", "no:randomly", *pytest_args]

            command_env = _subprocess_env()
            command_env["TERM"] = "xterm-256color"
            command_env["COLUMNS"] = str(
                term_columns + 1 if platform.system() == "Windows" else term_columns
            )
            command_env["PY_COLORS"] = "1"

            if stdin:
                # makes Console.is_terminal == True
                command_env["FORCE_COLOR"] = "true"

            command_env.update(env)

            with change_file(
                tmp_path / "conftest.py", lambda text: text + self.conftest_footer()
            ):
                result = sp.run(
                    cmd, cwd=tmp_path, capture_output=True, env=command_env, input=stdin
                )

            result_stdout = result.stdout.decode("utf-8")
            result_stderr = result.stderr.decode("utf-8")
            plain_stdout = normalize(result_stdout)

            console.print("run>", *cmd)

            console.print(Panel(Text.from_ansi(result_stdout), title="stdout"))
            if result_stderr:
                console.print(Panel(Text(result_stderr), title="stderr"))

            assert result.returncode == snapshot_arg(returncode)

            if stdout is not None:
                assert snapshot_arg(stdout) == result_stdout

            original = normalize(result_stderr).splitlines()
            lines = [
                line
                for line in original
                if not any(
                    s in line
                    for s in [
                        'No entry for terminal type "unknown"',
                        "using dumb terminal settings.",
                    ]
                )
            ]

            assert "\n".join(lines) == snapshot_arg(stderr)

            if report is not None:

                report_list = []
                record = False
                for line in plain_stdout.splitlines():
                    line = line.strip()
                    if line.startswith("===="):
                        record = False

                    if record and line:
                        report_list.append(line)

                    if (
                        line.startswith(("-----", "═════"))
                        and "inline-snapshot" in line
                    ):
                        record = True

                report_str = "\n".join(report_list)

                assert report_str == report, repr(report_str)

            error_str = (
                "\n".join(
                    [
                        error_line
                        for line in plain_stdout.splitlines()
                        if (error_line := _pytest_error_line(line)) is not None
                    ]
                )
                + "\n"
            )
            if not error_str.strip():
                error_str = ""

            if sys.version_info >= (3, 11):
                # assert rewriting gets disabled by inline-snapshot for older python versions
                assert snapshot_arg(error) == normalize(error_str)

            assert snapshot_arg(changed_files) == self._changed_files(tmp_path)

            assert snapshot_arg(outcomes) == parse_outcomes(plain_stdout.splitlines())

            return self._new_example(self._read_files(tmp_path))._next_seed()
