# Copyright 2020 The Chromium Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

from recipe_engine import post_process
from google.protobuf import field_mask_pb2
from google.protobuf import json_format
from google.protobuf.struct_pb2 import Struct

from PB.go.chromium.org.luci.buildbucket.proto import build as build_pb2
from PB.go.chromium.org.luci.buildbucket.proto import builder as builder_pb2
from PB.go.chromium.org.luci.buildbucket.proto import builds_service as builds_service_pb2
from PB.go.chromium.org.luci.buildbucket.proto import common as common_pb2

from PB.recipes.dart.roller.lkgr import LastKnownGoodRevision

DEPS = [
    'depot_tools/bot_update',
    'depot_tools/gclient',
    'depot_tools/git',
    'recipe_engine/context',
    'recipe_engine/json',
    'recipe_engine/path',
    'recipe_engine/properties',
    'recipe_engine/raw_io',
    'recipe_engine/runtime',
    'recipe_engine/step',
]

PROPERTIES = LastKnownGoodRevision


def RunSteps(api, properties):
  builders = properties.builders
  assert builders, 'the recipe requires builders to be specified'
  ref = properties.ref
  assert ref, 'the recipe requires a ref to be specified'
  builds = _find_good_builds(api, builders)
  commit_hash = _find_good_hash(api, builders, builds)

  # Guard against the empty string to avoid deleting the remote branch
  if not commit_hash:
    step = api.step('no good hashes found', None)
    step.presentation.logs['builds'] = _pretty_print(api, builds)
    return

  api.gclient.set_config('dart')
  with api.context(cwd=api.path['cache'].join('builder')):
    api.bot_update.ensure_checkout(timeout=1080, no_fetch_tags=True)

  git_args = ['push']
  if api.runtime.is_experimental:
    git_args.append('--dry-run')
  git_args += [
      'https://dart.googlesource.com/sdk.git',
      '%s:%s' % (commit_hash, ref),
  ]
  api.git(*git_args, name='push %s to %s' % (commit_hash, ref))


def _find_good_builds(api, builders):
  batch_request = builds_service_pb2.BatchRequest(
      requests=[dict(search_builds=_search_req(api, b)) for b in builders])
  request_dict = json_format.MessageToDict(batch_request)
  try:
    api.step(
        'search for good builds',
        ['bb', 'batch'],
        infra_step=True,
        stdin=api.json.input(request_dict),
        stdout=api.json.output(name='response'),
        timeout=5 * 60,  # 5 minutes
    )
  finally:
    step_result = api.step.active_result
    step_result.presentation.logs['request'] = _pretty_print(api, request_dict)

    batch_result = builds_service_pb2.BatchResponse()
    json_format.ParseDict(
        step_result.stdout or {},
        batch_result,
        # Do not fail the build because recipe's proto copy is stale.
        ignore_unknown_fields=True)

    builds = {}
    step_text = []
    for i, result in enumerate(batch_result.responses):
      # Print response errors in step text.
      if result.HasField('error'):
        step_text.extend([
            'Request #%d' % i,
            'Status code: %s' % result.error.code,
            'Message: %s' % result.error.message,
            '',  # Blank line.
        ])
      elif result.search_builds.builds:
        commit = lambda build: build.output.properties['got_revision']
        commit_hashes = map(commit, result.search_builds.builds)
        builder = result.search_builds.builds[0].builder.builder
        builds[builder] = commit_hashes
    step_result.presentation.step_text = '<br>'.join(step_text)

  return builds


def _search_req(api, builder):
  return builds_service_pb2.SearchBuildsRequest(
      predicate=builds_service_pb2.BuildPredicate(
          builder=builder,
          status=common_pb2.SUCCESS,
          tags=[
              # Ignore non-standard builds (e.g. bisections)
              common_pb2.StringPair(key='user_agent', value='luci-scheduler')
          ],
      ),
      fields=field_mask_pb2.FieldMask(
          paths=['builds.*.output.properties', 'builds.*.builder']))


def _find_good_hash(api, builders, builds):
  first = builders[0]
  rest = builders[1:]

  # The builds are sorted newest-to-oldest.
  candidates = builds.get(first.builder, [])
  if not candidates:
    return None
  intersection = set(candidates)
  for builder in rest:
    intersection.intersection_update(builds.get(builder.builder, []))
    if not intersection:
      return None
  for commit_hash in candidates:
    if commit_hash in intersection:
      return commit_hash

  assert False, 'A commit hash must be in the intersection.'  # pragma: no cover


def _pretty_print(api, the_dict):
  return api.json.dumps(the_dict, indent=2, sort_keys=True).splitlines()


def GenTests(api):
  yield api.test(
      'no-builders-no-steps',
      api.post_process(post_process.DoesNotRunRE, '(search|push).*'),
      api.expect_exception('AssertionError'),
      api.post_process(post_process.StatusException),
      api.post_process(post_process.DropExpectation),
  )

  builders = [
      builder_pb2.BuilderID(project='dart', bucket='ci', builder=builder)
      for builder in ['a', 'b', 'c']
  ]
  input_properties = api.properties(
      LastKnownGoodRevision(builders=builders, ref='refs/heads/lkgr'))

  def _search_response(builds):
    return dict(
        search_builds=builds_service_pb2.SearchBuildsResponse(builds=builds))

  no_builds_response = json_format.MessageToDict(
      builds_service_pb2.BatchResponse(responses=[_search_response([])]))
  yield api.test(
      'no-builds',
      input_properties,
      api.step_data(
          'search for good builds', stdout=api.json.output(no_builds_response)),
      api.post_process(post_process.MustRun, 'no good hashes found'),
      api.post_process(post_process.DoesNotRunRE, 'push.*'),
      api.post_process(post_process.StatusSuccess),
  )

  yield api.test(
      'error',
      input_properties,
      api.step_data(
          'search for good builds',
          retcode=1,
          stdout=api.json.output(
              json_format.MessageToDict(
                  builds_service_pb2.BatchResponse(responses=[
                      dict(error=dict(code=1, message='search failed'))
                  ])))),
      api.post_process(post_process.DoesNotRunRE, '(no|push).*'),
      api.post_process(post_process.StatusException),
  )

  def build(builder, got_revision):
    properties = Struct()
    properties.update({'got_revision': got_revision})
    return build_pb2.Build(
        builder=builder,
        output=build_pb2.Build.Output(properties=properties),
    )

  responses = []
  responses.append(_search_response([build(builders[0], '2')]))
  responses.append(
      _search_response(
          [build(builders[1], commit) for commit in ['1', '2', '3']]))
  responses.append(
      _search_response([build(builders[2], commit) for commit in ['2', '3']]))
  test_response = builds_service_pb2.BatchResponse(responses=responses)
  search_data = api.step_data(
      'search for good builds',
      stdout=api.json.output(json_format.MessageToDict(test_response)))
  yield api.test(
      'push',
      input_properties,
      search_data,
      api.post_process(post_process.MustRun, 'push 2 to refs/heads/lkgr'),
      api.post_process(post_process.DoesNotRunRE, 'no good hashes found'),
      api.post_process(post_process.StatusSuccess),
  )
  yield api.test(
      'dry-run',
      input_properties,
      search_data,
      api.runtime(is_luci=True, is_experimental=True),
      api.post_process(post_process.MustRun, 'push 2 to refs/heads/lkgr'),
      api.post_process(post_process.DoesNotRunRE, 'no good hashes found'),
      api.post_process(post_process.StepCommandContains,
                       'push 2 to refs/heads/lkgr', ['--dry-run']),
      api.post_process(post_process.StatusSuccess),
      api.post_process(post_process.DropExpectation),
  )

  no_intersection = _search_response(
      [build(builders[0], commit) for commit in ['0', '1', '4']])
  no_intersection_response = json_format.MessageToDict(
      builds_service_pb2.BatchResponse(
          responses=[no_intersection, responses[1], responses[2]]))
  yield api.test(
      'no-intersection',
      input_properties,
      api.step_data(
          'search for good builds',
          stdout=api.json.output(no_intersection_response)),
      api.post_process(post_process.MustRun, 'no good hashes found'),
      api.post_process(post_process.DoesNotRunRE, 'push.*'),
      api.post_process(post_process.StatusSuccess),
  )