# Copyright 2020 The Pigweed Authors
#
# Licensed under the Apache License, Version 2.0 (the "License"); you may not
# use this file except in compliance with the License. You may obtain a copy of
# the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
# License for the specific language governing permissions and limitations under
# the License.
"""Update a CIPD package to the latest version."""

import collections
import dataclasses
import re
from typing import Generator, List

from PB.recipe_engine import result as result_pb2
from PB.go.chromium.org.luci.buildbucket.proto import common as common_pb2
from PB.recipes.pigweed.cipd_roller import InputProperties, Package
from recipe_engine import recipe_test_api

DEPS = [
    'fuchsia/auto_roller',
    'pigweed/checkout',
    'recipe_engine/cipd',
    'recipe_engine/file',
    'recipe_engine/json',
    'recipe_engine/path',
    'recipe_engine/properties',
    'recipe_engine/step',
]

PROPERTIES = InputProperties


def _is_platform(part):
    """Return true for platform-style strings.

    Example matches: "linux-amd64", "${platform}", "${os}-amd64", "cp38", etc.
    Example non-matches: "clang", "amd64", "linux".
    """
    if '{' in part:
        return True

    # Match Python version indicators.
    if re.match(r'cp\d+', part):
        return True

    try:
        os, arch = part.split('-')
        return os in ('linux', 'mac', 'windows')
    except ValueError:
        return False


def find_shared_tags(api, package_tags, tag):
    """Attempts to find a tag shared by all packages.

    This function can be used if the intersection of the sets of tags
    associated with different-platform packages with the same 'ref' is empty.
    It finds a tag shared by all packages, with as many of them as possible
    matching 'ref'.
    """
    # Find the most common tags.  We use the sorted dict keys for determinism.
    package_paths = sorted(package_tags.keys())
    counter = collections.Counter()
    for path in package_paths:
        counter.update(package_tags[path])
    most_common_tags = counter.most_common()

    with api.step.nest("find shared tag"):
        for tag_candidate, _ in most_common_tags:
            # There is at least one package for which the version with the
            # specified 'ref' does not have this tag. See if there exists a
            # version of this package that *does* have this tag.  If so, use
            # that version.
            updated_tags = dict()
            for package_path in package_paths:
                if tag_candidate in package_tags[package_path]:
                    # For this package we already have a version with this tag,
                    # nothing to do.
                    continue
                try:
                    package_data = api.cipd.describe(
                        package_path, tag_candidate
                    )
                except api.step.StepFailure:
                    # No luck: there exists no version with this tag.
                    break
                updated_tags[package_path] = set(
                    x.tag
                    for x in package_data.tags
                    if x.tag.startswith(tag + ':')
                )
            else:
                # We found a version of each package with the tag_candidate.
                merged_tags = dict()
                merged_tags.update(package_tags)
                merged_tags.update(updated_tags)
                tags = set.intersection(*merged_tags.values())
                # Should always succeed.
                assert len(tags) > 0
                # Update package_tags to be consistent with the returned tags.
                package_tags.update(updated_tags)
                return tags

    # We failed to find any tag that meets our criteria.
    return set()


_FOOTER = """
CQ-Do-Not-Cancel-Tryjobs: true
Build-Errors: continue
""".strip()


@dataclasses.dataclass(order=True)
class Roll:
    package_name: str
    old_version: str
    new_version: str

    def message(self):
        return f'From: {self.old_version}\nTo: {self.new_version}'


@dataclasses.dataclass
class Commit:
    rolls: List[Roll] = dataclasses.field(default_factory=list)

    def message(self):
        rolls = sorted(self.rolls)

        result = []
        result.append(f'roll: {", ".join(x.package_name for x in rolls)}')

        if len(rolls) == 1:
            result.append('')
            result.append(rolls[0].message())

        else:
            for roll in rolls:
                result.append('')
                result.append(roll.package_name)
                result.append(roll.message())

        return '\n'.join(result)

    def __bool__(self):
        return bool(self.rolls)


def process_package(api, checkout, pkg):
    json_path = checkout.root.joinpath(*re.split(r'[\\/]+', pkg.json_path))

    if not pkg.name:
        # Turn foo/bar/baz/${platform} and foo/bar/baz/${os=mac}-${arch}
        # into 'baz'.
        pkg.name = [
            part for part in pkg.spec.split('/') if not _is_platform(part)
        ][-1]

    basename = api.path.basename(json_path)
    cipd_json = api.file.read_json(f'read {basename}', json_path)
    packages = cipd_json
    if isinstance(cipd_json, dict):
        packages = cipd_json['packages']
    old_version = None
    package = None
    for package in packages:
        if package['path'] == pkg.spec:
            old_version = package['tags'][0]
            break
    else:
        raise api.step.StepFailure(
            f"couldn't find package {pkg.spec} in {json_path}"
        )

    assert package.get('platforms'), 'platforms empty in json'
    platforms = package.get('platforms')
    base, name = pkg.spec.rstrip('/').rsplit('/', 1)
    if _is_platform(name):
        package_paths = [f'{base}/{x}' for x in platforms]
    else:
        package_paths = [pkg.spec]

    package_tags = {}
    tags = None
    for package_path in package_paths:
        try:
            package_data = api.cipd.describe(package_path, pkg.ref)

        except api.step.StepFailure:
            # If we got here this package doesn't have the correct ref. This
            # is likely because it's a new platform for an existing package.
            # In that case ignore this platform when checking that refs
            # agree on package versions. We still need at least one platform
            # to have the ref or the checks below will fail.
            pass

        else:
            package_tags[package_path] = set(
                x.tag
                for x in package_data.tags
                if x.tag.startswith(pkg.tag + ':')
            )
            if tags is None:
                tags = set(package_tags[package_path])
            else:
                tags.intersection_update(package_tags[package_path])

    if not tags and pkg.allow_mismatched_refs:
        # The package with the requested ref has non-overlapping tag values
        # for different platforms.  Try relaxing the requirement that all
        # packages come from the same ref, and see if this allows us to find
        # a set with shared tag values.
        tags = find_shared_tags(api, package_tags, pkg.tag)

    with api.step.nest('common tags') as presentation:
        presentation.step_summary_text = '\n'.join(sorted(tags))

    if not tags:
        err_lines = [f'no common tags across "{pkg.ref}" refs of packages']
        for package_path, package_tags in sorted(package_tags.items()):
            err_lines.append('')
            err_lines.append(package_path)
            for tag in package_tags:
                err_lines.append(tag)

        raise api.step.StepFailure('<br>'.join(err_lines))

    # Deterministically pick one of the common tags.
    new_version = sorted(tags)[0]
    package['tags'] = [new_version]

    version_part = new_version.split(':', 1)[1]
    match = re.search(r'(?:\d|\b)(rc|pre|beta|alpha)(?:\d|\b)', version_part)
    if match:
        raise api.step.StepFailure(
            f'found pre-release indicator {match.group(1)!r} in '
            f'{version_part!r}'
        )

    # Verify there's only one instance of each platform package with this
    # tag.
    with api.step.nest('check number of instances'):
        for package_path in package_paths:
            api.cipd.describe(package_path, new_version)

    with api.step.nest('new_version') as presentation:
        presentation.step_summary_text = new_version

    if old_version in tags:
        with api.step.nest('already up-to-date') as presentation:
            presentation.step_summary_text = (
                'current version {} in common tags'
            ).format(old_version)
        return None

    else:
        api.file.write_text(
            f'write {basename}',
            json_path,
            api.json.dumps(cipd_json, indent=2, separators=(',', ': ')) + '\n',
        )
        return Roll(
            package_name=pkg.name,
            old_version=old_version,
            new_version=new_version,
        )


def RunSteps(api, props):
    props.checkout_options.use_trigger = False
    checkout = api.checkout(props.checkout_options)

    commit = Commit()

    for pkg in props.packages:
        with api.step.nest(pkg.spec):
            roll = process_package(api, checkout, pkg)
        if roll:
            commit.rolls.append(roll)

    if not commit:
        return result_pb2.RawResult(
            summary_markdown='nothing to roll',
            status=common_pb2.SUCCESS,
        )

    change = api.auto_roller.attempt_roll(
        props.auto_roller_options,
        repo_dir=checkout.root,
        commit_message=commit.message(),
    )

    return api.auto_roller.raw_result(change)


def GenTests(api) -> Generator[recipe_test_api.TestData, None, None]:
    """Create tests."""

    def package(cipd_path, old_version, platforms=None):
        result = {
            'path': cipd_path,
            'tags': [old_version],
            '_comment': 'comments should be preserved',
        }

        if platforms is not None:
            result['platforms'] = list(platforms)

        return result

    def describe(spec, package_path, tags, **kwargs):
        return api.step_data(
            f'{spec}.cipd describe {package_path}',
            api.cipd.example_describe(
                package_path,
                test_data_tags=[f'{name}:{value}' for name, value in tags],
            ),
            **kwargs,
        )

    def no_ref(spec, package_paths):
        res = None
        for package_path in package_paths:
            step = describe(spec, package_path, (), retcode=1)
            res = res + step if res else step
        return res

    def no_common_tags(spec, package_paths, tagname='git_revision'):
        res = None
        for i, package_path in enumerate(package_paths):
            step = describe(spec, package_path, ((tagname, i),))
            res = res + step if res else step
        return res

    def multiple_common_tags(
        spec,
        package_paths,
        tagname='git_revision',
        versions=(1, 2, 3),
    ):
        res = None
        for package_path in package_paths:
            step = describe(
                spec,
                package_path,
                tuple((tagname, x) for x in versions),
            )
            res = res + step if res else step
        return res

    def relaxing_works(spec, package_paths, tagname='git_revision'):
        return api.step_data(
            f'{spec}.find shared tag.cipd describe {package_paths[1]}',
            api.cipd.example_describe(
                package_paths[1],
                test_data_tags=[f'{tagname}:{0}'],
            ),
        )

    def relaxing_does_not_work(spec, package_paths, tagname='git_revision'):
        # No version of package_paths[1] with hoped-for tag.
        return api.step_data(
            f'{spec}.find shared tag.cipd describe {package_paths[0]}',
            api.cipd.example_describe(package_paths[0]),
            retcode=1,
        ) + api.step_data(
            f'{spec}.find shared tag.cipd describe {package_paths[1]}',
            api.cipd.example_describe(package_paths[1]),
            retcode=1,
        )

    def read_file_step_data(spec, package_json_name, *packages):
        return api.step_data(
            f'{spec}.read {package_json_name}',
            api.file.read_json({'packages': packages}),
        )

    prefix = 'pigweed/host_tools/cp38'
    spec = f'{prefix}/${{platform}}'
    paths = (
        f'{prefix}/linux-amd64',
        f'{prefix}/windows-amd64',
    )

    def package_props(**kwargs):
        kwargs.setdefault('spec', spec)
        kwargs.setdefault('ref', 'latest')
        kwargs.setdefault('tag', 'git_revision')
        kwargs.setdefault(
            'json_path',
            'pw_env_setup/py/pw_env_setup/cipd_setup/pigweed.json',
        )
        return Package(**kwargs)

    def properties(packages, dry_run=True, **kwargs):
        props = InputProperties(**kwargs)
        props.packages.extend(packages)
        props.checkout_options.CopyFrom(api.checkout.git_options())
        props.auto_roller_options.CopyFrom(
            api.auto_roller.Options(
                dry_run=dry_run,
                remote=props.checkout_options.remote,
            )
        )
        return api.properties(props)

    yield (
        api.test('success')
        + properties([package_props()])
        + api.checkout.ci_test_data()
        + read_file_step_data(
            spec,
            'pigweed.json',
            package(
                spec,
                'git_revision:123',
                platforms=['linux-amd64', 'windows-amd64'],
            ),
        )
        + api.auto_roller.dry_run_success()
    )

    yield (
        api.test('rc', status='FAILURE')
        + properties([package_props(tag='version')])
        + api.checkout.ci_test_data()
        + read_file_step_data(
            spec,
            'pigweed.json',
            package(
                spec,
                'version:123',
                platforms=['linux-amd64', 'windows-amd64'],
            ),
        )
        + describe(spec, paths[0], (('version', '234-rc3'),))
        + describe(spec, paths[1], (('version', '234-rc3'),))
    )

    yield (
        api.test('multiple')
        + properties(
            [
                package_props(spec='foo/${platform}'),
                package_props(spec='bar/${platform}'),
            ]
        )
        + api.checkout.ci_test_data()
        + read_file_step_data(
            'foo/${platform}',
            'pigweed.json',
            package(
                'foo/${platform}',
                'git_revision:foo123',
                platforms=['linux-amd64', 'windows-amd64'],
            ),
            package(
                'bar/${platform}',
                'git_revision:bar123',
                platforms=['linux-amd64', 'windows-amd64'],
            ),
        )
        + read_file_step_data(
            'bar/${platform}',
            'pigweed.json',
            package(
                'foo/${platform}',
                'git_revision:397a2597cdc237f3026e6143b683be4b9ab60540',
                platforms=['linux-amd64', 'windows-amd64'],
            ),
            package(
                'bar/${platform}',
                'git_revision:bar123',
                platforms=['linux-amd64', 'windows-amd64'],
            ),
        )
        + api.auto_roller.dry_run_success()
    )

    bad_spec = f'bad-{spec}'
    yield (
        api.test('bad_package_spec', status='FAILURE')
        + properties([package_props(spec=bad_spec)])
        + api.checkout.ci_test_data()
        + read_file_step_data(
            bad_spec, 'pigweed.json', package(spec, 'git_revision:123')
        )
    )

    yield (
        api.test('no_common_tags', status='FAILURE')
        + properties([package_props()])
        + api.checkout.ci_test_data()
        + read_file_step_data(
            spec,
            'pigweed.json',
            package(
                spec,
                'git_revision:123',
                platforms=('linux-amd64', 'windows-amd64'),
            ),
        )
        + no_common_tags(spec, paths)
    )

    yield (
        api.test('no_common_tags_but_relaxing_ref_mismatch_helps')
        + properties([package_props(allow_mismatched_refs=True)])
        + api.checkout.ci_test_data()
        + read_file_step_data(
            spec,
            'pigweed.json',
            package(
                spec,
                'git_revision:123',
                platforms=('linux-amd64', 'windows-amd64'),
            ),
        )
        # Package 0 exists at git_revision:0, package 1 at git_revision:1
        + no_common_tags(spec, paths)
        # However, package 1 also exists at git_revision:0, so that revision is
        # good to use.
        + relaxing_works(spec, paths)
        + api.auto_roller.dry_run_success()
    )

    yield (
        api.test(
            'no_common_tags_and_relaxing_ref_mismatch_does_not_help',
            status='FAILURE',
        )
        + properties([package_props(allow_mismatched_refs=True)])
        + api.checkout.ci_test_data()
        + read_file_step_data(
            spec,
            'pigweed.json',
            package(
                spec,
                'git_revision:123',
                platforms=('linux-amd64', 'windows-amd64'),
            ),
        )
        # Package 0 exists at git_revision:0, package 1 at git_revision:1
        + no_common_tags(spec, paths)
        # However, package 1 does not exist at git_revision:0.
        + relaxing_does_not_work(spec, paths)
    )

    yield (
        api.test('multiple_common_tags')
        + properties([package_props()])
        + api.checkout.ci_test_data()
        + multiple_common_tags(spec, paths)
        + read_file_step_data(
            spec,
            'pigweed.json',
            package(
                spec,
                'git_revision:2',
                platforms=('linux-amd64', 'windows-amd64'),
            ),
        )
    )

    yield (
        api.test('missing_tag')
        + properties([package_props()])
        + api.checkout.ci_test_data()
        + multiple_common_tags(spec, paths)
        + no_ref(spec, [f'{prefix}/fake-amd64'])
        + read_file_step_data(
            spec,
            'pigweed.json',
            package(
                spec,
                'git_revision:2',
                platforms=('linux-amd64', 'windows-amd64', 'fake-amd64'),
            ),
        )
    )

    no_curly_spec = 'pigweed/host_tools/linux-amd64'
    yield (
        api.test('no_curlies_in_spec')
        + properties([package_props(spec=no_curly_spec)])
        + api.checkout.ci_test_data()
        + read_file_step_data(
            no_curly_spec,
            'pigweed.json',
            package(
                no_curly_spec,
                'git_revision:123',
                platforms=('linux-amd64', 'windows-amd64'),
            ),
        )
        + api.auto_roller.dry_run_success()
    )

    yield (
        api.test('platform-independent')
        + properties([package_props(spec='foo/bar/baz')])
        + api.checkout.ci_test_data()
        + read_file_step_data(
            'foo/bar/baz',
            'pigweed.json',
            package(
                'foo/bar/baz',
                'git_revision:123',
                platforms=['linux-amd64', 'mac-amd64'],
            ),
        )
        + api.auto_roller.dry_run_success()
    )