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

from __future__ import annotations

import json
from textwrap import dedent

import pytest

from pants.backend.python.goals.lockfile import (
    GeneratePythonLockfile,
    RequestedPythonUserResolveNames,
    setup_user_lockfile_requests,
)
from pants.backend.python.goals.lockfile import rules as lockfile_rules
from pants.backend.python.subsystems.setup import RESOLVE_OPTION_KEY__DEFAULT, PythonSetup
from pants.backend.python.target_types import PythonRequirementTarget
from pants.backend.python.util_rules import pex
from pants.backend.python.util_rules.interpreter_constraints import InterpreterConstraints
from pants.core.goals.generate_lockfiles import GenerateLockfileResult, UserGenerateLockfiles
from pants.engine.fs import DigestContents
from pants.engine.internals.scheduler import ExecutionError
from pants.testutil.python_rule_runner import PythonRuleRunner
from pants.testutil.rule_runner import PYTHON_BOOTSTRAP_ENV, QueryRule
from pants.util.ordered_set import FrozenOrderedSet
from pants.util.strutil import strip_prefix


@pytest.fixture
def rule_runner() -> PythonRuleRunner:
    rule_runner = PythonRuleRunner(
        rules=[
            *lockfile_rules(),
            *pex.rules(),
            QueryRule(GenerateLockfileResult, [GeneratePythonLockfile]),
        ]
    )
    rule_runner.set_options([], env_inherit=PYTHON_BOOTSTRAP_ENV)
    return rule_runner


def _generate_and_check_metadata(
    *,
    rule_runner: PythonRuleRunner,
    separate_metadata_file: bool = False,
    requirements_strings: list[str] | None = None,
    no_binary: bool = False,
    only_binary: bool = False,
    requirement_constraints_str: str = '//   "requirement_constraints": [],\n',
    excludes_and_overrides_str: str = '//   "excludes": [],\n//   "overrides": [],\n',
    sources_str: str = '//   "sources": []',
) -> str:
    if requirements_strings is None:
        requirements_strings = ["ansicolors==1.1.8"]
    result = rule_runner.request(
        GenerateLockfileResult,
        [
            GeneratePythonLockfile(
                requirements=FrozenOrderedSet(requirements_strings),
                find_links=FrozenOrderedSet([]),
                interpreter_constraints=InterpreterConstraints(),
                resolve_name="test",
                lockfile_dest="test.lock",
                diff=False,
                lock_style="universal",
                complete_platforms=(),
            )
        ],
    )
    digest_contents = rule_runner.request(DigestContents, [result.digest])
    if separate_metadata_file:
        assert len(digest_contents) == 2
        metadata = json.loads(digest_contents[1].content.decode())
    else:
        assert len(digest_contents) == 1
        metadata = None
    content = digest_contents[0].content.decode()

    requirements_strings.sort()
    requirements_header_formatting = ",\n".join(
        f'            //     "{req}"' for req in requirements_strings
    )
    pex_header = (
        dedent(
            f"""\
            // This lockfile was autogenerated by Pants. To regenerate, run:
            //
            //    pants generate-lockfiles --resolve=test
            //
            // --- BEGIN PANTS LOCKFILE METADATA: DO NOT EDIT OR REMOVE ---
            // {{
            //   "version": 6,
            //   "valid_for_interpreter_constraints": [],
            //   "generated_with_requirements": [
{requirements_header_formatting}
            //   ],
            //   "manylinux": "manylinux2014",
            """
        )
        + requirement_constraints_str
        + (
            dedent(
                """\
            //   "only_binary": [
            //     "ansicolors"
            //   ],
            """
            )
            if only_binary
            else '//   "only_binary": [],\n'
        )
        + (
            dedent(
                """\
            //   "no_binary": [
            //     "ansicolors"
            //   ],
            """
            )
            if no_binary
            else '//   "no_binary": [],\n'
        )
        + excludes_and_overrides_str
        + sources_str
        + dedent(
            """\
            ,
            //   "lock_style": "universal",
            //   "complete_platforms": []
            // }
            // --- END PANTS LOCKFILE METADATA ---
            """
        )
    )

    if metadata:
        assert metadata == {
            "description": "This lockfile was generated by Pants. To regenerate, run: pants generate-lockfiles --resolve=test",
            "version": 6,
            "valid_for_interpreter_constraints": [],
            "generated_with_requirements": requirements_strings,
            "manylinux": "manylinux2014",
            "excludes": [],
            "no_binary": ["ansicolors"] if no_binary else [],
            "only_binary": ["ansicolors"] if only_binary else [],
            "overrides": [],
            "sources": [],
            "requirement_constraints": [],
            "lock_style": "universal",
            "complete_platforms": [],
        }
        return content
    else:
        assert content.startswith(pex_header)
        return strip_prefix(content, pex_header)


def _generate(rule_runner: PythonRuleRunner, **kwargs) -> str:
    # Set defaults for new required fields
    kwargs.setdefault("lock_style", "universal")
    kwargs.setdefault("complete_platforms", ())

    result = rule_runner.request(
        GenerateLockfileResult,
        [
            GeneratePythonLockfile(
                resolve_name="test",
                lockfile_dest="test.lock",
                diff=False,
                **kwargs,
            )
        ],
    )
    digest_contents = rule_runner.request(DigestContents, [result.digest])
    content = digest_contents[0].content.decode()

    return content


@pytest.mark.parametrize(
    ("no_binary", "only_binary", "separate_metadata_file"),
    (
        (False, False, False),
        (False, True, False),
        (True, False, False),
        (True, False, True),
    ),
)
def test_pex_lockfile_generation(
    rule_runner: PythonRuleRunner,
    no_binary: bool,
    only_binary: bool,
    separate_metadata_file: bool,
) -> None:
    args = ["--python-resolves={'test': 'foo.lock'}"]
    no_binary_arg = f"{{'{RESOLVE_OPTION_KEY__DEFAULT}': ['ansicolors']}}"
    if no_binary:
        args.append(f"--python-resolves-to-no-binary={no_binary_arg}")
    if only_binary:
        args.append(f"--python-resolves-to-only-binary={no_binary_arg}")
    if separate_metadata_file:
        args.append("--python-separate-lockfile-metadata-file")
    rule_runner.set_options(args, env_inherit=PYTHON_BOOTSTRAP_ENV)

    lock_entry = json.loads(
        _generate_and_check_metadata(
            rule_runner=rule_runner,
            separate_metadata_file=separate_metadata_file,
            no_binary=no_binary,
            only_binary=only_binary,
        )
    )
    reqs = lock_entry["locked_resolves"][0]["locked_requirements"]
    assert len(reqs) == 1
    assert reqs[0]["project_name"] == "ansicolors"
    assert reqs[0]["version"] == "1.1.8"

    wheel = {
        "algorithm": "sha256",
        "hash": "00d2dde5a675579325902536738dd27e4fac1fd68f773fe36c21044eb559e187",
        "url": (
            "https://files.pythonhosted.org/packages/53/18/"
            + "a56e2fe47b259bb52201093a3a9d4a32014f9d85071ad07e9d60600890ca/"
            + "ansicolors-1.1.8-py2.py3-none-any.whl"
        ),
    }
    sdist = {
        "algorithm": "sha256",
        "hash": "99f94f5e3348a0bcd43c82e5fc4414013ccc19d70bd939ad71e0133ce9c372e0",
        "url": (
            "https://files.pythonhosted.org/packages/76/31/"
            + "7faed52088732704523c259e24c26ce6f2f33fbeff2ff59274560c27628e/"
            + "ansicolors-1.1.8.zip"
        ),
    }

    artifacts = reqs[0]["artifacts"]

    if not no_binary and not only_binary:
        # Don't assume that the order in artifacts is deterministic.
        # We can't just convert to a set because dicts aren't hashable.
        assert len(artifacts) == 2
        assert wheel in artifacts
        assert sdist in artifacts
    elif no_binary:
        assert artifacts == [sdist]
    elif only_binary:
        assert artifacts == [wheel]


def test_constraints_file(rule_runner: PythonRuleRunner) -> None:
    rule_runner.write_files({"constraints.txt": "ansicolors==1.1.7"})
    rule_runner.set_options(
        [
            "--python-resolves={'test': 'foo.lock'}",
            f"--python-resolves-to-constraints-file={{'{RESOLVE_OPTION_KEY__DEFAULT}': 'constraints.txt'}}",
        ],
        env_inherit=PYTHON_BOOTSTRAP_ENV,
    )

    lock_entry = json.loads(
        _generate_and_check_metadata(
            rule_runner=rule_runner,
            requirements_strings=["ansicolors>=1.0"],
            requirement_constraints_str=dedent(
                """\
                //   "requirement_constraints": [
                //     "ansicolors==1.1.7"
                //   ],
                """
            ),
        )
    )
    reqs = lock_entry["locked_resolves"][0]["locked_requirements"]
    assert len(reqs) == 1
    assert reqs[0]["project_name"] == "ansicolors"
    assert reqs[0]["version"] == "1.1.7"


def test_multiple_resolves() -> None:
    rule_runner = PythonRuleRunner(
        rules=[
            setup_user_lockfile_requests,
            *PythonSetup.rules(),
            QueryRule(UserGenerateLockfiles, [RequestedPythonUserResolveNames]),
        ],
        target_types=[PythonRequirementTarget],
    )
    rule_runner.write_files(
        {
            "BUILD": dedent(
                """\
                python_requirement(
                    name='a',
                    requirements=['a'],
                    resolve='a',
                )
                python_requirement(
                    name='b',
                    requirements=['b'],
                    resolve='b',
                )
                """
            ),
        }
    )
    rule_runner.set_options(
        [
            "--python-resolves={'a': 'a.lock', 'b': 'b.lock'}",
            # Override interpreter constraints for 'b', but use default for 'a'.
            "--python-resolves-to-interpreter-constraints={'b': ['==3.9.*']}",
            "--python-enable-resolves",
        ],
        env_inherit=PYTHON_BOOTSTRAP_ENV,
    )
    result = rule_runner.request(
        UserGenerateLockfiles, [RequestedPythonUserResolveNames(["a", "b"])]
    )
    assert set(result) == {
        GeneratePythonLockfile(
            requirements=FrozenOrderedSet(["a"]),
            find_links=FrozenOrderedSet([]),
            interpreter_constraints=InterpreterConstraints(["CPython>=3.9,<3.15"]),
            resolve_name="a",
            lockfile_dest="a.lock",
            diff=False,
            lock_style="universal",
            complete_platforms=(),
        ),
        GeneratePythonLockfile(
            requirements=FrozenOrderedSet(["b"]),
            find_links=FrozenOrderedSet([]),
            interpreter_constraints=InterpreterConstraints(["==3.9.*"]),
            resolve_name="b",
            lockfile_dest="b.lock",
            diff=False,
            lock_style="universal",
            complete_platforms=(),
        ),
    }


def test_empty_requirements(rule_runner: PythonRuleRunner) -> None:
    with pytest.raises(ExecutionError) as excinfo:
        json.loads(
            _generate_and_check_metadata(
                rule_runner=rule_runner,
                requirements_strings=[],
            )
        )

    assert (
        "Cannot generate lockfile with no requirements. Please add some requirements to test."
        in str(excinfo.value)
    )


def test_define_source_from_different_index(rule_runner: PythonRuleRunner) -> None:
    args = [
        "--python-resolves={'test': 'foo.lock'}",
        "--python-repos-indexes=['https://pypi.org/simple/','test_pipy=https://test.pypi.org/simple/']",
        "--python-resolves-to-sources={'test': ['test_pipy=cowsay']}",
    ]
    rule_runner.set_options(args, env_inherit=PYTHON_BOOTSTRAP_ENV)

    lock_entry = json.loads(
        _generate_and_check_metadata(
            rule_runner=rule_runner,
            requirements_strings=["cowsay==6.0", "ansicolors==1.1.8"],
            sources_str=dedent(
                """\
                //   "sources": [
                //     "test_pipy=cowsay"
                //   ]"""
            ),
        )
    )

    reqs = lock_entry["locked_resolves"][0]["locked_requirements"]

    assert len(reqs) == 2

    # Ansicolors is from PyPI (https://files.pythonhosted.org)
    assert reqs[0]["project_name"] == "ansicolors"
    assert reqs[0]["version"] == "1.1.8"

    wheel_ansicolors = {
        "algorithm": "sha256",
        "hash": "00d2dde5a675579325902536738dd27e4fac1fd68f773fe36c21044eb559e187",
        "url": (
            "https://files.pythonhosted.org/packages/53/18/"
            + "a56e2fe47b259bb52201093a3a9d4a32014f9d85071ad07e9d60600890ca/"
            + "ansicolors-1.1.8-py2.py3-none-any.whl"
        ),
    }
    sdist_ansicolors = {
        "algorithm": "sha256",
        "hash": "99f94f5e3348a0bcd43c82e5fc4414013ccc19d70bd939ad71e0133ce9c372e0",
        "url": (
            "https://files.pythonhosted.org/packages/76/31/"
            + "7faed52088732704523c259e24c26ce6f2f33fbeff2ff59274560c27628e/"
            + "ansicolors-1.1.8.zip"
        ),
    }
    artifacts = reqs[0]["artifacts"]
    assert wheel_ansicolors in artifacts
    assert sdist_ansicolors in artifacts

    # Cowsay is from TestPyPI (https://test-files.pythonhosted.org)
    assert reqs[1]["project_name"] == "cowsay"
    assert reqs[1]["version"] == "6.0"

    wheel_cowsay = {
        "algorithm": "sha256",
        "hash": "77b07c508af48aa300a90f3b3c5c013a12360a71fc5c87b1efb763fe2803a775",
        "url": (
            "https://test-files.pythonhosted.org/packages/73/56/"
            + "7922bfc226ccd44221befb6b866aaa4da4170bc9d8b036ef0675ce0f1b53/"
            + "cowsay-6.0-py2.py3-none-any.whl"
        ),
    }
    sdist_cowsay = {
        "algorithm": "sha256",
        "hash": "47445cb273684618a1786db8e8d05ec9258455f7eb74893e5d0933daafeb44ba",
        "url": (
            "https://test-files.pythonhosted.org/packages/7e/b5/"
            + "e8e802ddc7f5219417dc7d7953eec81ffe48ad129b793f3040361e4aff89/"
            + "cowsay-6.0.tar.gz"
        ),
    }
    artifacts = reqs[1]["artifacts"]
    assert wheel_cowsay in artifacts
    assert sdist_cowsay in artifacts


def test_override_version(rule_runner: PythonRuleRunner) -> None:
    args = [
        "--python-resolves={'test': 'test.lock'}",
        "--python-separate-lockfile-metadata-file",
        "--python-resolves-to-overrides={'test': ['python-dateutil==2.6.0']}",
    ]

    rule_runner.set_options(args, env_inherit=PYTHON_BOOTSTRAP_ENV)

    lock_entry = json.loads(
        _generate(
            rule_runner=rule_runner,
            requirements=FrozenOrderedSet(["freezegun==1.2.1"]),
            interpreter_constraints=InterpreterConstraints(),
            find_links=FrozenOrderedSet(),
        )
    )

    reqs = lock_entry["locked_resolves"][0]["locked_requirements"]
    print(reqs)
    assert reqs[0]["project_name"] == "freezegun"
    assert reqs[1]["project_name"] == "python-dateutil"
    # Without Override: python-dateutil>=2.7
    assert reqs[1]["version"] == "2.6.0"


def test_excluded_by_excludes(rule_runner: PythonRuleRunner) -> None:
    args = [
        "--python-resolves={'test': 'test.lock'}",
        "--python-separate-lockfile-metadata-file",
        "--python-resolves-to-excludes={'test': ['six']}",
    ]

    rule_runner.set_options(args, env_inherit=PYTHON_BOOTSTRAP_ENV)

    lock_entry = json.loads(
        _generate(
            rule_runner=rule_runner,
            requirements=FrozenOrderedSet(["fasteners==0.16.3"]),
            interpreter_constraints=InterpreterConstraints(),
            find_links=FrozenOrderedSet(),
        )
    )

    reqs = lock_entry["locked_resolves"][0]["locked_requirements"]
    print(reqs)

    assert reqs[0]["project_name"] == "fasteners"
    # Kept in in the metadata
    assert "six" in reqs[0]["requires_dists"]
    # But excluded as a project
    assert "six" not in {req["project_name"]: req for req in reqs}
