# Copyright 2019 The Fuchsia 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 google.protobuf import text_format as textpb
from recipe_engine import recipe_api

import base64
import os

BUILD_INIT_VERSION = 'git_revision:bd277978791e1e35bf76e0df57f6aef14aef0757'

DEFAULT_SPEC_DIR = 'infra/config/generated'


class SpecApi(recipe_api.RecipeApi):
  """API for working with textproto files."""

  def __init__(self, *args, **kwargs):
    super(SpecApi, self).__init__(*args, **kwargs)
    self._build_init_path = None

  class ParseError(Exception):
    pass

  class NotFoundError(Exception):
    pass

  def _ensure_build_init(self):
    with self.m.step.nest('ensure build_init'):
      with self.m.context(infra_steps=True):
        pkgs = self.m.cipd.EnsureFile()
        pkgs.add_package('fuchsia/infra/build_init/${platform}',
                         BUILD_INIT_VERSION)
        build_init_dir = self.m.path['start_dir'].join('cipd', 'build_init')
        self.m.cipd.ensure(build_init_dir, pkgs)
        self._build_init_path = build_init_dir.join('build_init')
        return self._build_init_path

  def load(self,
           spec_remote,
           Type,
           spec_dir=DEFAULT_SPEC_DIR,
           spec_revision=None):
    spec, _ = self.get_spec_revision(spec_remote, Type, spec_dir, spec_revision)
    return spec

  def get_spec_revision(self,
                        spec_remote,
                        Type,
                        spec_dir=DEFAULT_SPEC_DIR,
                        spec_revision=None):
    """Loads a spec message from the given Buildbucket build input.

    If the build input is for a gitiles commit, the Gitiles host, project,
    and ref are inferred from the build input and the spec is fetched from
    Gitiles.

    If the build input is for a gerrit change, the project is rebased on top
    of the parent gitiles_commit in build.

    Args:
      spec_remote (string): URL of the specs git repository.
      spec_dir (string): The directory within spec_remote containing specs.
        Specifically we'll look for the spec in
        "<spec_dir>/<project>/specs/<bucket>/<builder>.textproto".
      spec_revision (string or None): The git revision to fetch the spec from.

    Returns:
      spec (Type), rev (string): The parsed proto API spec and its revision.

    Raises:
      NotFoundError if the spec file could not be located.
      ParseError if the proto fails to parse into the given Type.
    """
    build = self.m.buildbucket.build
    with self.m.context(infra_steps=True):
      self._ensure_build_init()
      assert self._build_init_path
      # build_init takes a text build.proto on stdin. base64 encode just in case
      # the input contains non utf-8 text which breaks logdog.
      build_msg = base64.b64encode(build.SerializeToString())
      builder_name = build.builder.builder
      # Have subbuilders use the same specs as the parent builders.
      if builder_name.endswith('-subbuild'):
        builder_name = builder_name[:builder_name.index('-subbuild')]
      # This path is relative to the git remote, and doesn't contain any local
      # directory placeholder, so it's safe to use os.path.join() rather than
      # the standard recipe idiom of <Path object>.join().
      spec_path = os.path.join(spec_dir, build.builder.project, 'specs',
                               build.builder.bucket,
                               builder_name + '.textproto')
      if not spec_revision:
        assert build.input.gitiles_commit
        commit_remote = 'https://%s/%s' % (build.input.gitiles_commit.host,
                                           build.input.gitiles_commit.project)
        if commit_remote == spec_remote:
          # build_init will ignore the spec_revision in this case and use the
          # revision of the gitiles_commit associated with the build. This is just
          # for consistency to return the correct spec revision being used.
          spec_revision = build.input.gitiles_commit.id
        else:
          # Get revision at HEAD and send to build_init.
          ref = 'refs/heads/master'
          spec_revision = self.m.gitiles.refs(
              spec_remote,
              refspath='refs/heads',
              test_data=lambda: self.m.json.test_api.output(
                  {'refs/heads/master': 'deadbeef'})).get(ref, None)
          assert spec_revision, 'Failed to retrieve spec revision'
      cmd = [
          self._build_init_path, '-spec_remote', spec_remote, '-spec_path',
          spec_path, '-spec_ref', spec_revision
      ]

      # build_init creates a git checkout. Do this in separate directory to avoid
      # clashing with jiri checkout.
      with self.m.context(cwd=self.m.path.mkdtemp('build_init_workspace')):
        step_result = self.m.step(
            name='build_init',
            cmd=cmd,
            stdin=self.m.raw_io.input_text(data=build_msg),
            stdout=self.m.raw_io.output(),
            ok_ret=(0, 2),  # 0 = ok, 2 = file not found.
            timeout=10 * 60,
        )

      if step_result.retcode == 2:
        raise SpecApi.NotFoundError('could not find spec at "%s"' % spec_path)

      spec = step_result.stdout.strip()
      if not spec:
        # Return parse error if the file is empty. google.protobuf considers
        # that a valid proto so this results in an empty spec otherwise.
        raise SpecApi.ParseError('file is empty')

      try:
        return textpb.Merge(spec, Type()), spec_revision
      except Exception as e:
        raise SpecApi.ParseError(e)