from __future__ import annotations

import os
import re
import subprocess
import sys
import tempfile
from collections.abc import Iterator
from contextlib import contextmanager

import filelock

import mypy.api
from mypy.test.config import package_path, pip_lock, pip_timeout, test_temp_dir
from mypy.test.data import DataDrivenTestCase, DataSuite
from mypy.test.helpers import assert_string_arrays_equal, perform_file_operations

# NOTE: options.use_builtins_fixtures should not be set in these
# tests, otherwise mypy will ignore installed third-party packages.


class PEP561Suite(DataSuite):
    files = ["pep561.test"]
    base_path = "."

    def run_case(self, testcase: DataDrivenTestCase) -> None:
        test_pep561(testcase)


@contextmanager
def virtualenv(python_executable: str = sys.executable) -> Iterator[tuple[str, str]]:
    """Context manager that creates a virtualenv in a temporary directory

    Returns the path to the created Python executable
    """
    with tempfile.TemporaryDirectory() as venv_dir:
        proc = subprocess.run(
            [python_executable, "-m", "venv", venv_dir], cwd=os.getcwd(), capture_output=True
        )
        if proc.returncode != 0:
            err = proc.stdout.decode("utf-8") + proc.stderr.decode("utf-8")
            raise Exception("Failed to create venv.\n" + err)
        if sys.platform == "win32":
            yield venv_dir, os.path.abspath(os.path.join(venv_dir, "Scripts", "python"))
        else:
            yield venv_dir, os.path.abspath(os.path.join(venv_dir, "bin", "python"))


def upgrade_pip(python_executable: str) -> None:
    """Install pip>=21.3.1. Required for editable installs with PEP 660."""
    if (
        sys.version_info >= (3, 11)
        or (3, 10, 3) <= sys.version_info < (3, 11)
        or (3, 9, 11) <= sys.version_info < (3, 10)
    ):
        # Skip for more recent Python releases which come with pip>=21.3.1
        # out of the box - for performance reasons.
        return

    install_cmd = [python_executable, "-m", "pip", "install", "pip>=21.3.1"]
    try:
        with filelock.FileLock(pip_lock, timeout=pip_timeout):
            proc = subprocess.run(install_cmd, capture_output=True, env=os.environ)
    except filelock.Timeout as err:
        raise Exception(f"Failed to acquire {pip_lock}") from err
    if proc.returncode != 0:
        raise Exception(proc.stdout.decode("utf-8") + proc.stderr.decode("utf-8"))


def install_package(
    pkg: str, python_executable: str = sys.executable, editable: bool = False
) -> None:
    """Install a package from test-data/packages/pkg/"""
    working_dir = os.path.join(package_path, pkg)
    with tempfile.TemporaryDirectory() as dir:
        install_cmd = [python_executable, "-m", "pip", "install"]
        if editable:
            install_cmd.append("-e")
        install_cmd.append(".")

        # Note that newer versions of pip (21.3+) don't
        # follow this env variable, but this is for compatibility
        env = {"PIP_BUILD": dir}
        # Inherit environment for Windows
        env.update(os.environ)
        try:
            with filelock.FileLock(pip_lock, timeout=pip_timeout):
                proc = subprocess.run(install_cmd, cwd=working_dir, capture_output=True, env=env)
        except filelock.Timeout as err:
            raise Exception(f"Failed to acquire {pip_lock}") from err
    if proc.returncode != 0:
        raise Exception(proc.stdout.decode("utf-8") + proc.stderr.decode("utf-8"))


def test_pep561(testcase: DataDrivenTestCase) -> None:
    """Test running mypy on files that depend on PEP 561 packages."""
    assert testcase.old_cwd is not None, "test was not properly set up"
    python = sys.executable

    assert python is not None, "Should be impossible"
    pkgs, pip_args = parse_pkgs(testcase.input[0])
    mypy_args = parse_mypy_args(testcase.input[1])
    editable = False
    for arg in pip_args:
        if arg == "editable":
            editable = True
        else:
            raise ValueError(f"Unknown pip argument: {arg}")
    assert pkgs, "No packages to install for PEP 561 test?"
    with virtualenv(python) as venv:
        venv_dir, python_executable = venv
        if editable:
            # Editable installs with PEP 660 require pip>=21.3
            upgrade_pip(python_executable)
        for pkg in pkgs:
            install_package(pkg, python_executable, editable)

        cmd_line = list(mypy_args)
        has_program = not ("-p" in cmd_line or "--package" in cmd_line)
        if has_program:
            program = testcase.name + ".py"
            with open(program, "w", encoding="utf-8") as f:
                for s in testcase.input:
                    f.write(f"{s}\n")
            cmd_line.append(program)

        cmd_line.extend(["--no-error-summary", "--hide-error-codes"])
        if python_executable != sys.executable:
            cmd_line.append(f"--python-executable={python_executable}")

        steps = testcase.find_steps()
        if steps != [[]]:
            steps = [[]] + steps

        for i, operations in enumerate(steps):
            perform_file_operations(operations)

            output = []
            # Type check the module
            out, err, returncode = mypy.api.run(cmd_line)

            # split lines, remove newlines, and remove directory of test case
            for line in (out + err).splitlines():
                if line.startswith(test_temp_dir + os.sep):
                    output.append(line[len(test_temp_dir + os.sep) :].rstrip("\r\n"))
                else:
                    # Normalize paths so that the output is the same on Windows and Linux/macOS.
                    # Yes, this is naive: replace all slashes preceding first colon, if any.
                    path, *rest = line.split(":", maxsplit=1)
                    if rest:
                        path = path.replace(os.sep, "/")
                    output.append(":".join([path, *rest]).rstrip("\r\n"))
            iter_count = "" if i == 0 else f" on iteration {i + 1}"
            expected = testcase.output if i == 0 else testcase.output2.get(i + 1, [])

            assert_string_arrays_equal(
                expected,
                output,
                f"Invalid output ({testcase.file}, line {testcase.line}){iter_count}",
            )

        if has_program:
            os.remove(program)


def parse_pkgs(comment: str) -> tuple[list[str], list[str]]:
    if not comment.startswith("# pkgs:"):
        return ([], [])
    else:
        pkgs_str, *args = comment[7:].split(";")
        return ([pkg.strip() for pkg in pkgs_str.split(",")], [arg.strip() for arg in args])


def parse_mypy_args(line: str) -> list[str]:
    m = re.match("# flags: (.*)$", line)
    if not m:
        return []  # No args; mypy will spit out an error.
    return m.group(1).split()