#!/usr/bin/env python
# Copyright 2022 The IREE Authors
#
# Licensed under the Apache License v2.0 with LLVM Exceptions.
# See https://llvm.org/LICENSE.txt for license information.
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

# Pushes module patches to an appropriate repository.
#
# Typical steps:
# 1. Advance the main branches on our forks to get all upstream commits
#    (if you forget to do this, your cherry-pick will likely complain about
#    a bad ref).
# 2. In your submodule:
#      git fetch
#      git cherry-pick <some commit>
# 3. Run this script from the main IREE repository (one of the following):
#      patch_module.py --module=llvm-project
#      patch_module.py --module=stablehlo
# 4. Send a PR on the main IREE repo to bump the submodule. Be sure to include
#    the name of the patch branch for posterity.

import argparse
from datetime import date
import os
import sys

import iree_utils
import iree_modules

PATCH_REMOTE_ALIAS = "patched"


def main(args):
  module_info = iree_modules.MODULE_INFOS.get(args.module)
  if not module_info:
    raise SystemExit(f"ERROR: Bad value for --module. Must be one of: "
                     f"{', '.join(iree_modules.MODULE_INFOS.keys())}")

  if args.command == "patch":
    main_patch(args, module_info)
  else:
    raise SystemExit(
        f"ERROR: Unrecognized --command. Must be one of: patch, unpatch")


def main_patch(args, module_info: iree_modules.ModuleInfo):
  module_root = os.path.join(iree_utils.get_repo_root(), module_info.path)
  setup_module_remotes(module_root, module_info)

  branch_name = find_unused_branch_name(module_info)
  print(f"Allocated branch: {branch_name}")
  current_commit, summary = iree_utils.git_current_commit(repo_dir=module_root)
  print(f"Module is currently at: {summary}")
  print(
      f"*** Pushing branch {branch_name} to {module_info.fork_repository_push} ***"
  )
  print(f"(Please ignore any messages below about creating a PR)\n")
  iree_utils.git_exec([
      "push", PATCH_REMOTE_ALIAS, f"{current_commit}:refs/heads/{branch_name}"
  ],
                      repo_dir=module_root)
  print(f"*** Branch {branch_name} pushed ***")

  print(f"******* Congratulations *******")
  print(
      f"You have pushed your commits to {branch_name} on {module_info.fork_repository_push}."
  )
  print(
      f"Your main repository should now show that the submodule has been edited."
  )
  print(f"Make a commit, referencing the above branch cherry-picks and ")
  print(f"land the resulting PR.")
  print(f"You can push more commits to this module's patch branch via:")
  print(
      f"  (cd {module_info.path} && git push {PATCH_REMOTE_ALIAS} HEAD:{branch_name})"
  )


def setup_module_remotes(module_root: str,
                         module_info: iree_modules.ModuleInfo):
  iree_utils.git_setup_remote(PATCH_REMOTE_ALIAS,
                              url=module_info.fork_repository_push,
                              repo_dir=module_root)


def find_unused_branch_name(module_info: iree_modules.ModuleInfo):
  branch_base = f"{module_info.branch_prefix}{date.today().strftime('%Y%m%d')}"
  branch_name = branch_base
  existing_branches = iree_utils.git_ls_remote_branches(
      module_info.fork_repository_pull,
      filter=[f"refs/heads/{module_info.branch_prefix}*"])
  i = 1
  while branch_name in existing_branches:
    branch_name = f"{branch_base}.{i}"
    i += 1
  return branch_name


def parse_arguments(argv):
  parser = argparse.ArgumentParser(description="IREE Submodule Patcher")
  parser.add_argument("--module", help="Submodule to operate on", default=None)
  parser.add_argument("--command", help="Command to execute", default="patch")
  args = parser.parse_args(argv)
  return args


if __name__ == "__main__":
  main(parse_arguments(sys.argv[1:]))