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

from recipe_engine import recipe_api
from recipe_engine.recipe_test_api import StepTestData
from typing import Callable, Optional


class GerritApi(recipe_api.RecipeApi):
  """Module for interact with Gerrit endpoints"""

  def __init__(self, *args, **kwargs):
    super(GerritApi, self).__init__(*args, **kwargs)
    self._changes_target_branch_cache = {}

  def __call__(self, name, cmd, infra_step=True, **kwargs):
    """Wrapper for easy calling of gerrit_utils steps."""
    assert isinstance(cmd, (list, tuple))
    prefix = "gerrit "

    env = self.m.context.env
    env.setdefault("PATH", "%(PATH)s")
    env["PATH"] = self.m.path.pathsep.join(
      [env["PATH"], str(self.repo_resource())]
    )

    with self.m.context(env=env):
      return self.m.step(
        prefix + name,
        ["vpython3", self.repo_resource("gerrit_client.py")] + cmd,
        infra_step=infra_step,
        **kwargs,
      )

  def call_raw_api(
    self,
    host,
    path,
    method=None,
    body=None,
    accept_statuses=None,
    name=None,
    verbose=False,
    **kwargs,
  ):
    """Call an arbitrary Gerrit API that returns a JSON response.

    Returns:
      The JSON response data.
    """
    args = [
      "rawapi",
      "--host",
      host,
      "--path",
      path,
      "--json_file",
      self.m.json.output(),
    ]
    if method:
      args.extend(["--method", method])
    if body:
      args.extend(["--body", self.m.json.dumps(body)])
    if accept_statuses:
      args.extend(
        ["--accept_status", ",".join(str(i) for i in accept_statuses)]
      )
    if verbose:
      args.append("--verbose")

    step_name = name or "call_raw_api (%s)" % path
    step_result = self(step_name, args, **kwargs)
    return step_result.json.output

  def create_gerrit_branch(
    self, host, project, branch, commit, verbose=False, **kwargs
  ):
    """Creates a new branch from given project and commit

    Returns:
      The ref of the branch created
    """
    args = [
      "branch",
      "--host",
      host,
      "--project",
      project,
      "--branch",
      branch,
      "--commit",
      commit,
      "--json_file",
      self.m.json.output(),
    ]
    allow_existent_branch = kwargs.pop("allow_existent_branch", False)
    if allow_existent_branch:
      args.append("--allow-existent-branch")
    if verbose:
      args.append("--verbose")
    step_name = "create_gerrit_branch (%s %s)" % (project, branch)
    step_result = self(step_name, args, **kwargs)
    ref = step_result.json.output.get("ref")
    return ref

  def create_gerrit_tag(
    self, host, project, tag, commit, verbose=False, **kwargs
  ):
    """Creates a new tag at the given commit.

    Returns:
      The ref of the tag created.
    """
    args = [
      "tag",
      "--host",
      host,
      "--project",
      project,
      "--tag",
      tag,
      "--commit",
      commit,
      "--json_file",
      self.m.json.output(),
    ]
    if verbose:
      args.append("--verbose")
    step_name = "create_gerrit_tag (%s %s)" % (project, tag)
    step_result = self(step_name, args, **kwargs)
    ref = step_result.json.output.get("ref")
    return ref

  # TODO(machenbach): Rename to get_revision? And maybe above to
  # create_ref?
  def get_gerrit_branch(self, host, project, branch, verbose=False, **kwargs):
    """Gets a branch from given project and commit

    Returns:
      The revision of the branch
    """
    args = [
      "branchinfo",
      "--host",
      host,
      "--project",
      project,
      "--branch",
      branch,
      "--json_file",
      self.m.json.output(),
    ]
    if verbose:
      args.append("--verbose")
    step_name = "get_gerrit_branch (%s %s)" % (project, branch)
    step_result = self(step_name, args, **kwargs)
    revision = step_result.json.output.get("revision")
    return revision

  def get_change_description(
    self,
    host,
    change,
    patchset,
    timeout=None,
    step_test_data=None,
    verbose=True,
  ):
    """Gets the description for a given CL and patchset.

    Args:
      host: URL of Gerrit host to query.
      change: The change number.
      patchset: The patchset number.
      verbose: Whether to enable verbose logging.

    Returns:
      The description corresponding to given CL and patchset.
    """
    ri = self.get_revision_info(
      host, change, patchset, timeout, step_test_data, verbose
    )
    return ri["commit"]["message"]

  def get_revision_info(
    self,
    host,
    change,
    patchset,
    timeout=None,
    step_test_data=None,
    verbose=True,
  ):
    """
    Returns the info for a given patchset of a given change.

    Args:
      host: Gerrit host to query.
      change: The change number.
      patchset: The patchset number.
      verbose: Whether to enable verbose logging.

    Returns:
      A dict for the target revision as documented here:
          https://gerrit-review.googlesource.com/Documentation/rest-api-changes.html#list-changes
    """
    assert int(change), change
    assert int(patchset), patchset

    step_test_data = step_test_data or (
      lambda: self.test_api.get_one_change_response_data(
        change_number=change, patchset=patchset
      )
    )

    cls = self.get_changes(
      host,
      query_params=[("change", str(change))],
      o_params=["ALL_REVISIONS", "ALL_COMMITS"],
      limit=1,
      timeout=timeout,
      step_test_data=step_test_data,
      verbose=verbose,
    )
    cl = cls[0] if len(cls) == 1 else {"revisions": {}}
    for ri in cl["revisions"].values():
      # TODO(tandrii): add support for patchset=='current'.
      if str(ri["_number"]) == str(patchset):
        return ri

    raise self.m.step.InfraFailure(
      "Error querying for CL description: host:%r change:%r; patchset:%r"
      % (host, change, patchset)
    )

  def _is_us_east_coast(self):
    bot_dims = self.m.buildbucket.swarming_bot_dimensions
    if not bot_dims:
      return False
    for dim in bot_dims:
      if dim.key == "zone" and dim.value in ("us-atl", "us-iad"):
        return True
    return False

  def get_changes(
    self,
    host,
    query_params,
    start=None,
    limit=None,
    o_params=None,
    step_test_data=None,
    verbose=True,
    **kwargs,
  ):
    """Queries changes for the given host.

    Args:
      * host: URL of Gerrit host to query.
      * query_params: Query parameters as list of (key, value) tuples to form a
          query as documented here:
          https://gerrit-review.googlesource.com/Documentation/user-search.html#search-operators
      * start: How many changes to skip (starting with the most recent).
      * limit: Maximum number of results to return.
      * o_params: A list of additional output specifiers, as documented here:
          https://gerrit-review.googlesource.com/Documentation/rest-api-changes.html#list-changes
      * step_test_data: Optional mock test data for the underlying gerrit client.
      * verbose: Whether to enable verbose logging.

    Returns:
      A list of change dicts as documented here:
          https://gerrit-review.googlesource.com/Documentation/rest-api-changes.html#list-changes
    """
    if not step_test_data:
      step_test_data = lambda: self.test_api.get_one_change_response_data()  # noqa: E731

    step_name = kwargs.pop("name", "changes")

    def _run_query(target_host):
      args = [
        "changes",
        "--host",
        target_host,
        "--json_file",
        self.m.json.output(),
      ]
      if start:
        args += ["--start", str(start)]
      if limit:
        args += ["--limit", str(limit)]
      for k, v in query_params:
        args += ["-p", "%s=%s" % (k, v)]
      for v in o_params or []:
        args += ["-o", v]
      if verbose:
        args.append("--verbose")

      return self(
        step_name, args, step_test_data=step_test_data, **kwargs
      ).json.output

    original_res = _run_query(host)

    # We've seen different results be returned for the same list-changes request
    # depending on the location of the requester. So we retry against multiple
    # mirrors in that case.
    # TODO(crbug.com/530541275): Remove this if Gerrit replication issues are
    # resolved.
    host_mirrors = {
      "https://chromium-review.googlesource.com": [
        "https://us1-mirror-chromium-review.googlesource.com",
        "https://us2-mirror-chromium-review.googlesource.com",
      ],
    }
    if not original_res and self._is_us_east_coast():
      for mirror in host_mirrors.get(host, []):
        res = _run_query(mirror)
        if res:
          return res

    return original_res

  def get_file_content(
    self,
    host,
    change,
    path,
    revision="current",
    step_test_data=None,
    verbose=False,
  ):
    """Gets the content of a file from a change.

    Args:
      host: Gerrit host to query.
      change: The change number.
      path: The file path.
      revision: The revision number or ID.
      step_test_data: Optional mock test data.
      verbose: Whether to enable verbose logging.

    Returns:
      The content of the file.
    """
    args = [
      "content",
      "--host",
      host,
      "--change",
      change,
      "--path",
      path,
      "--revision",
      revision,
    ]
    if verbose:
      args.append("--verbose")

    if not step_test_data:
      step_test_data = lambda: self.m.raw_io.test_api.stream_output_text(  # noqa: E731
        "example content"
      )

    return self(
      "get file content",
      args,
      stdout=self.m.raw_io.output_text(),
      step_test_data=step_test_data,
    ).stdout

  def get_related_changes(
    self, host, change, revision="current", step_test_data=None, verbose=False
  ):
    """Queries related changes for a given host, change, and revision.

    Args:
      * host: URL of Gerrit host to query.
      * change: The change-id of the change to get related changes for as
          documented here:
          https://gerrit-review.googlesource.com/Documentation/rest-api-changes.html#change-id
      * revision: The revision-id of the revision to get related changes for as
          documented here:
          https://gerrit-review.googlesource.com/Documentation/rest-api-changes.html#revision-id
          This defaults to current, which names the most recent patch set.
      * step_test_data: Optional mock test data for the underlying gerrit client.
      * verbose: Whether to enable verbose logging.

    Returns:
      A related changes dictionary as documented here:
          https://gerrit-review.googlesource.com/Documentation/rest-api-changes.html#related-changes-info

    """
    args = [
      "relatedchanges",
      "--host",
      host,
      "--change",
      change,
      "--revision",
      revision,
      "--json_file",
      self.m.json.output(),
    ]
    if verbose:
      args.append("--verbose")
    if not step_test_data:
      step_test_data = lambda: self.test_api.get_related_changes_response_data()  # noqa: E731

    return self(
      "relatedchanges", args, step_test_data=step_test_data
    ).json.output

  def abandon_change(
    self,
    host,
    change,
    message=None,
    name=None,
    step_test_data=None,
    verbose=False,
  ):
    args = [
      "abandon",
      "--host",
      host,
      "--change",
      int(change),
      "--json_file",
      self.m.json.output(),
    ]
    if message:
      args.extend(["--message", message])
    if verbose:
      args.append("--verbose")
    if not step_test_data:
      step_test_data = lambda: self.test_api.get_one_change_response_data(  # noqa: E731
        status="ABANDONED", _number=str(change)
      )

    return self(
      name or "abandon",
      args,
      step_test_data=step_test_data,
    ).json.output

  def restore_change(
    self,
    host,
    change,
    message=None,
    name=None,
    step_test_data=None,
    verbose=False,
  ):
    args = [
      "restore",
      "--host",
      host,
      "--change",
      change,
      "--json_file",
      self.m.json.output(),
    ]
    if message:
      args.extend(("--message", message))
    if verbose:
      args.append("--verbose")
    if not step_test_data:
      step_test_data = lambda: self.test_api.get_one_change_response_data(  # noqa: E731
        status="NEW", _number=str(change)
      )

    return self(
      name or "restore",
      args,
      step_test_data=step_test_data,
    ).json.output

  def set_change_label(
    self,
    host,
    change,
    label_name,
    label_value,
    name=None,
    step_test_data=None,
    verbose=False,
  ):
    args = [
      "setlabel",
      "--host",
      host,
      "--change",
      int(change),
      "--json_file",
      self.m.json.output(),
      "-l",
      label_name,
      label_value,
    ]
    if verbose:
      args.append("--verbose")
    return self(
      name or "setlabel",
      args,
      step_test_data=step_test_data,
    ).json.output

  def add_message(
    self,
    host: str,
    change: int,
    message: str,
    revision: str | int = "current",
    automatic_attention_set_update: Optional[bool] = None,
    step_name: str = None,
    step_test_data: Callable[[], StepTestData] | None = None,
    verbose: bool = False,
  ) -> None:
    """Add a message to a change at given revision.

    Args:
      * host: URL of Gerrit host of the change.
      * change: The ID of the change to add message to as documented here:
          https://gerrit-review.googlesource.com/Documentation/rest-api-changes.html#change-id
      * message: The content of the message to add to the change.
      * revision: The ID of the revision of change to add message to as
          documented here:
          https://gerrit-review.googlesource.com/Documentation/rest-api-changes.html#revision-id
          This defaults to current, which names the most recent patchset.
      * automatic_attention_set_update: Whether to update the attention set.
      * step_name: Optional step name.
      * step_test_data: Optional mock test data for the underlying gerrit
          client.
      * verbose: Whether to enable verbose logging.
    """
    args = [
      "addmessage",
      "--host",
      host,
      "--change",
      int(change),
      "--revision",
      str(revision),
      "--message",
      message,
      "--json_file",
      self.m.json.output(),
    ]
    if automatic_attention_set_update is not None:
      args += [
        "--automatic-attention-set-update"
        if automatic_attention_set_update
        else "--no-automatic-attention-set-update"
      ]
    if verbose:
      args.append("--verbose")
    if not step_test_data:
      step_test_data = lambda: self.m.json.test_api.output({})  # noqa: E731
    return self(
      step_name or "add message",
      args,
      step_test_data=step_test_data,
    ).json.output

  def move_changes(
    self,
    host,
    project,
    from_branch,
    to_branch,
    step_test_data=None,
    verbose=False,
  ):
    args = [
      "movechanges",
      "--host",
      host,
      "-p",
      "project=%s" % project,
      "-p",
      "branch=%s" % from_branch,
      "-p",
      "status=open",
      "--destination_branch",
      to_branch,
      "--json_file",
      self.m.json.output(),
    ]
    if verbose:
      args.append("--verbose")

    if not step_test_data:
      step_test_data = lambda: self.test_api.get_one_change_response_data(  # noqa: E731
        branch=to_branch
      )

    return self(
      "move changes",
      args,
      step_test_data=step_test_data,
    ).json.output

  def update_files(
    self,
    host,
    project,
    branch,
    new_contents_by_file_path,
    commit_msg,
    params=frozenset(["status=NEW"]),
    cc_list=frozenset([]),
    submit=False,
    submit_later=False,
    step_test_data_create_change=None,
    step_test_data_submit_change=None,
    verbose=False,
  ):
    """Update a set of files by creating and submitting a Gerrit CL.

    Args:
      * host: URL of Gerrit host to name.
      * project: Gerrit project name, e.g. chromium/src.
      * branch: The branch to land the change, e.g. main
      * new_contents_by_file_path: Dict of the new contents with file path as
          the key.
      * commit_msg: Description to add to the CL.
      * params: A list of additional ChangeInput specifiers, with format
          'key=value'.
      * cc_list: A list of addresses to notify.
      * submit: Should land this CL instantly.
      * submit_later: If this change has related CLs, we may want to commit
           them in a chain. So only set Bot-Commit+1, making it ready for
           submit together. Ignored if submit is True.
      * step_test_data_create_change: Optional mock test data for the step
           create gerrit change.
      * step_test_data_submit_change: Optional mock test data for the step
          submit gerrit change.
      * verbose: Whether to enable verbose logging.


    Returns:
      A ChangeInfo dictionary as documented here:
          https://gerrit-review.googlesource.com/Documentation/rest-api-changes.html#create-change
          Or if the change is submitted, here:
          https://gerrit-review.googlesource.com/Documentation/rest-api-changes.html#submit-change
    """
    assert len(new_contents_by_file_path) > 0, (
      "The dict of file paths should not be empty."
    )
    create_command = [
      "createchange",
      "--host",
      host,
      "--project",
      project,
      "--branch",
      branch,
      "--subject",
      commit_msg,
      "--json_file",
      self.m.json.output(),
    ]
    for p in params:
      create_command.extend(["-p", p])
    for cc in cc_list:
      create_command.extend(["--cc", cc])
    if verbose:
      create_command.append("--verbose")
    step_test_data = step_test_data_create_change or (
      lambda: self.test_api.update_files_response_data()
    )

    step_result = self(
      "create change at (%s %s)" % (project, branch),
      create_command,
      step_test_data=step_test_data,
    )
    change = int(step_result.json.output.get("_number"))
    step_result.presentation.links["change %d" % change] = "%s/#/q/%d" % (
      host,
      change,
    )

    with self.m.step.nest("update contents in CL %d" % change):
      for path, content in new_contents_by_file_path.items():
        _file = self.m.path.mkstemp()
        self.m.file.write_raw(
          "store the new content for %s" % path, _file, content
        )
        edit_file_command = [
          "changeedit",
          "--host",
          host,
          "--change",
          change,
          "--path",
          path,
          "--file",
          _file,
        ]
        if verbose:
          edit_file_command.append("--verbose")
        self("edit file %s" % path, edit_file_command)

    publish_command = [
      "publishchangeedit",
      "--host",
      host,
      "--change",
      change,
    ]
    if verbose:
      publish_command.append("--verbose")
    self("publish edit", publish_command)

    NEW_PATCHSET_REVISION_NUMBER = 2
    # Make sure the new patchset is propagated to Gerrit backend.
    with self.m.step.nest("verify the patchset exists on CL %d" % change):
      retries = 0
      max_retries = 2
      while retries <= max_retries:
        try:
          if self.get_revision_info(
            host, change, patchset=NEW_PATCHSET_REVISION_NUMBER
          ):
            break
        except self.m.step.InfraFailure:
          if retries == max_retries:  # pragma: no cover
            raise
          retries += 1
          with self.m.step.nest("waiting before retry"):
            self.m.time.sleep((2**retries) * 10)

    if submit or submit_later:
      set_bot_commit_command = [
        "setbotcommit",
        "--host",
        host,
        "--change",
        change,
        "--revision",
        NEW_PATCHSET_REVISION_NUMBER,
      ]
      if verbose:
        set_bot_commit_command.append("--verbose")
      self("set Bot-Commit+1 for change %d" % change, set_bot_commit_command)
    if submit:
      submit_command = [
        "submitchange",
        "--host",
        host,
        "--change",
        change,
        "--json_file",
        self.m.json.output(),
      ]
      if verbose:
        submit_command.append("--verbose")
      step_test_data = step_test_data_submit_change or (
        lambda: self.test_api.update_files_response_data(status="MERGED")
      )
      step_result = self(
        "submit change %d" % change,
        submit_command,
        step_test_data=step_test_data,
      )
    return step_result.json.output