# Copyright 2022 Pants project contributors (see CONTRIBUTORS.md).
# Licensed under the Apache License, Version 2.0 (see LICENSE).

import itertools
import os
import shutil
import subprocess

import pytest

from pex.cli.testing import run_pex3
from pex.common import safe_open
from pex.fetcher import URLFetcher
from pex.resolve.locked_resolve import Artifact, FileArtifact, LockedRequirement
from pex.resolve.lockfile import json_codec
from pex.resolve.lockfile.model import Lockfile
from pex.resolve.resolver_configuration import PipConfiguration
from pex.testing import make_env, run_pex_command
from pex.typing import TYPE_CHECKING

if TYPE_CHECKING:
    from typing import Any, Iterator


def iter_requirements(lf):
    # type: (Lockfile) -> Iterator[LockedRequirement]
    return itertools.chain.from_iterable(
        locked_resolve.locked_requirements for locked_resolve in lf.locked_resolves
    )


def assert_file_artifact(artifact):
    # type: (Artifact) -> FileArtifact
    assert isinstance(artifact, FileArtifact)
    return artifact


def iter_artifacts(lf):
    # type: (Lockfile) -> Iterator[FileArtifact]
    return itertools.chain.from_iterable(
        map(assert_file_artifact, locked_requirement.iter_artifacts())
        for locked_requirement in iter_requirements(lf)
    )


@pytest.fixture
def lock(tmpdir):
    # type: (Any) -> str
    return os.path.join(str(tmpdir), "lock")


@pytest.fixture
def lock_file(lock):
    # type: (str) -> Lockfile
    run_pex3(
        "lock", "create", "ansicolors==1.1.8", "--style", "universal", "-o", lock, "--indent", "2"
    ).assert_success()
    return json_codec.load(lock)


@pytest.fixture
def find_links(
    tmpdir,  # type: Any
    lock_file,  # type: Lockfile
):
    # type: (...) -> str

    url_fetcher = URLFetcher()
    find_links = os.path.join(str(tmpdir), "find_links")
    for artifact in iter_artifacts(lock_file):
        with url_fetcher.get_body_stream(artifact.url) as url_fp:
            with safe_open(os.path.join(find_links, artifact.filename), "wb") as fl_fp:
                shutil.copyfileobj(url_fp, fl_fp)

    # We need the current default --pip-version requirements for some tests that do PyPI offline
    # resolves.
    pip_version = PipConfiguration().version
    repository_pex = os.path.join(str(tmpdir), "repository.pex")
    run_pex_command(
        args=[
            pip_version.setuptools_requirement,
            pip_version.wheel_requirement,
            "--include-tools",
            "-o",
            repository_pex,
        ]
    ).assert_success()
    subprocess.check_call(
        args=[repository_pex, "repository", "extract", "-f", find_links], env=make_env(PEX_TOOLS=1)
    )
    return find_links


def test_lock_update_repo_migration_dry_run(
    lock,  # type: str
    lock_file,  # type: Lockfile
    find_links,  # type: str
):
    # type: (...) -> None

    result = run_pex3(
        "lock", "update", "--pin", "--no-pypi", "--find-links", find_links, lock, "--dry-run"
    )
    result.assert_success()
    assert (
        "Updates for lock generated by universal\n"
        "  Would update ansicolors 1.1.8 artifacts:\n"
        "    https://files.pythonhosted.org/packages/53/18/"
        "a56e2fe47b259bb52201093a3a9d4a32014f9d85071ad07e9d60600890ca/"
        "ansicolors-1.1.8-py2.py3-none-any.whl -> "
        "file://{find_links}/ansicolors-1.1.8-py2.py3-none-any.whl\n"
        "    https://files.pythonhosted.org/packages/76/31/"
        "7faed52088732704523c259e24c26ce6f2f33fbeff2ff59274560c27628e/ansicolors-1.1.8.zip -> "
        "file://{find_links}/ansicolors-1.1.8.zip\n".format(find_links=find_links)
    ) == result.output


def test_lock_update_repo_migration_dry_run_path_mapping(
    lock,  # type: str
    lock_file,  # type: Lockfile
    find_links,  # type: str
):
    # type: (...) -> None

    result = run_pex3(
        "lock",
        "update",
        "--pin",
        "--no-pypi",
        "--find-links",
        find_links,
        lock,
        "--dry-run",
        "--path-mapping",
        "FOO|{find_links}".format(find_links=find_links),
    )
    result.assert_success()
    assert (
        "Updates for lock generated by universal\n"
        "  Would update ansicolors 1.1.8 artifacts:\n"
        "    https://files.pythonhosted.org/packages/53/18/"
        "a56e2fe47b259bb52201093a3a9d4a32014f9d85071ad07e9d60600890ca/"
        "ansicolors-1.1.8-py2.py3-none-any.whl -> "
        "file://${FOO}/ansicolors-1.1.8-py2.py3-none-any.whl\n"
        "    https://files.pythonhosted.org/packages/76/31/"
        "7faed52088732704523c259e24c26ce6f2f33fbeff2ff59274560c27628e/ansicolors-1.1.8.zip -> "
        "file://${FOO}/ansicolors-1.1.8.zip\n"
    ) == result.output


def test_lock_update_repo_migration_nominal(
    lock,  # type: str
    lock_file,  # type: Lockfile
    find_links,  # type: str
):
    # type: (...) -> None

    artifacts_by_filename = {artifact.filename: artifact for artifact in iter_artifacts(lock_file)}

    run_pex3(
        "lock", "update", "--pin", "--no-pypi", "--find-links", find_links, lock
    ).assert_success()

    for updated_artifact in iter_artifacts(json_codec.load(lock)):
        original_artifact = artifacts_by_filename.pop(updated_artifact.filename)
        assert original_artifact.fingerprint == updated_artifact.fingerprint
        assert original_artifact.url != updated_artifact.url
        assert (
            "file://{}".format(os.path.join(find_links, original_artifact.filename))
            == updated_artifact.url
        )
    assert not artifacts_by_filename


def test_lock_update_repo_migration_corrupted(
    lock,  # type: str
    lock_file,  # type: Lockfile
    find_links,  # type: str
):
    # type: (...) -> None

    with open(os.path.join(find_links, "ansicolors-1.1.8.zip"), "ab") as fp:
        fp.write(b"changed")

    result = run_pex3(
        "lock",
        "update",
        "--pin",
        "--no-pypi",
        "--find-links",
        find_links,
        lock,
        "--dry-run",
        "--path-mapping",
        "FOO|{find_links}".format(find_links=find_links),
    )
    result.assert_failure()
    assert (
        "Updates for lock generated by universal\n"
        "  Would update ansicolors 1.1.8 artifacts:\n"
        "    https://files.pythonhosted.org/packages/53/18/"
        "a56e2fe47b259bb52201093a3a9d4a32014f9d85071ad07e9d60600890ca/"
        "ansicolors-1.1.8-py2.py3-none-any.whl -> "
        "file://${FOO}/ansicolors-1.1.8-py2.py3-none-any.whl\n"
        "    https://files.pythonhosted.org/packages/76/31/"
        "7faed52088732704523c259e24c26ce6f2f33fbeff2ff59274560c27628e/"
        "ansicolors-1.1.8.zip"
        "#sha256:99f94f5e3348a0bcd43c82e5fc4414013ccc19d70bd939ad71e0133ce9c372e0 -> "
        "file://${FOO}/ansicolors-1.1.8.zip"
        "#sha256:9c872ada674e45fe740ebd4d00267617875e76d9bec4b3ce2e194e83661680dc\n"
    ) == result.output
    assert (
        "Detected fingerprint changes in the following locked project for lock generated by "
        "universal!\n"
        "ansicolors 1.1.8\n"
    ) == result.error


def test_lock_update_repo_migration_artifacts_removed(
    lock,  # type: str
    lock_file,  # type: Lockfile
    find_links,  # type: str
):
    # type: (...) -> None

    os.unlink(os.path.join(find_links, "ansicolors-1.1.8-py2.py3-none-any.whl"))
    result = run_pex3(
        "lock",
        "update",
        "--pin",
        "--no-pypi",
        "--find-links",
        find_links,
        lock,
        "--dry-run",
        "--path-mapping",
        "FOO|{find_links}".format(find_links=find_links),
    )
    result.assert_success()
    assert (
        "Updates for lock generated by universal\n"
        "  Would update ansicolors 1.1.8 artifacts:\n"
        "    https://files.pythonhosted.org/packages/76/31/"
        "7faed52088732704523c259e24c26ce6f2f33fbeff2ff59274560c27628e/ansicolors-1.1.8.zip -> "
        "file://${FOO}/ansicolors-1.1.8.zip\n"
        "    - https://files.pythonhosted.org/packages/53/18/"
        "a56e2fe47b259bb52201093a3a9d4a32014f9d85071ad07e9d60600890ca/"
        "ansicolors-1.1.8-py2.py3-none-any.whl\n"
    ) == result.output


def test_lock_update_repo_migration_artifacts_added(
    lock,  # type: str
    lock_file,  # type: Lockfile
    find_links,  # type: str
):
    # type: (...) -> None

    shutil.copy(
        os.path.join(find_links, "ansicolors-1.1.8.zip"),
        os.path.join(find_links, "ansicolors-1.1.8.tar.gz"),
    )
    result = run_pex3(
        "lock",
        "update",
        "--pin",
        "--no-pypi",
        "--find-links",
        find_links,
        lock,
        "--dry-run",
        "--path-mapping",
        "FOO|{find_links}".format(find_links=find_links),
    )
    result.assert_success()
    assert (
        "Updates for lock generated by universal\n"
        "  Would update ansicolors 1.1.8 artifacts:\n"
        "    + file://${FOO}/ansicolors-1.1.8.tar.gz\n"
        "    https://files.pythonhosted.org/packages/53/18/"
        "a56e2fe47b259bb52201093a3a9d4a32014f9d85071ad07e9d60600890ca/"
        "ansicolors-1.1.8-py2.py3-none-any.whl -> "
        "file://${FOO}/ansicolors-1.1.8-py2.py3-none-any.whl\n"
        "    https://files.pythonhosted.org/packages/76/31/"
        "7faed52088732704523c259e24c26ce6f2f33fbeff2ff59274560c27628e/ansicolors-1.1.8.zip -> "
        "file://${FOO}/ansicolors-1.1.8.zip\n"
    ) == result.output, result.output
