#!/usr/bin/env python3
# Lint as: python3
# Copyright 2019 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

# pylint: disable=missing-docstring
"""submodule_versions.

Synchronizes the tracked SUBMODULE_VERSIONS.txt file with the submodule state
in git.

Typical usage:
--------------
Exporting current git submodule state to SUBMODULE_VERSIONS.txt:
  Syntax: ./scripts/git/submodule_versions.py export

Importing versions in SUBMODULE_VERSIONS.txt to git submodule state:
  Syntax: ./scripts/git/submodule_versions.py import

Checking whether SUBMODULE_VERSIONS.txt and git state are in sync:
  Syntax: ./scripts/git/submodule_versions.py check
"""

import argparse
import os
import re
import sys

import utils

VERSIONS_FILE = "SUBMODULE_VERSIONS.txt"


def get_submodule_versions(repo_dir):
  raw_status = utils.execute(["git", "submodule", "status"],
                             cwd=repo_dir,
                             silent=True,
                             capture_output=True).stdout
  status_lines = []
  for line in raw_status.splitlines():
    # Format is a status char followed by revision, space and path.
    m = re.match(r"""^.([0-9a-z]+)\s+([^\s]+)""", line)
    if m:
      # Output as just the commit hash followed by space and path.
      status_lines.append(m.group(1) + " " + m.group(2))
  return "\n".join(status_lines) + "\n"


def export_versions(repo_dir):
  current_versions = get_submodule_versions(repo_dir)
  versions_file_path = os.path.join(repo_dir, VERSIONS_FILE)
  print("*** Exporting current submodule versions to:", versions_file_path)
  with open(versions_file_path, "w", encoding="UTF-8") as f:
    f.write(current_versions)
  utils.execute(["git", "add", VERSIONS_FILE], cwd=repo_dir)


def parse_versions(versions_text):
  versions = dict()
  for line in versions_text.splitlines():
    comps = line.split(" ", maxsplit=2)
    if len(comps) != 2:
      continue
    versions[comps[1]] = comps[0]
  return versions


def get_diff_versions(repo_dir):
  current_versions = parse_versions(get_submodule_versions(repo_dir))
  with open(os.path.join(repo_dir, VERSIONS_FILE), "r", encoding="UTF-8") as f:
    written_versions = parse_versions(f.read())
  diff_versions = current_versions.items() ^ written_versions.items()
  return {
      k: (current_versions.get(k), written_versions.get(k))
      for k, _ in diff_versions
  }


def sync_and_update_submodules(repo_dir):
  print("*** Synchronizing/updating submodules")
  utils.execute(["git", "submodule", "sync"], cwd=repo_dir)
  utils.execute(["git", "submodule", "update"], cwd=repo_dir)


def import_versions(repo_dir):
  print("*** Importing versions to git submodule state")
  diff_versions = get_diff_versions(repo_dir)
  if not diff_versions:
    print("*** No submodule updates required")
    return
  for path, (current, written) in diff_versions.items():
    if current is None:
      print(("Warning: Submodule %s does not exist but is "
             "still in the version file") % (path,))
      continue
    if written is None:
      print(f"Warning: Submodule '{current}' is not in the version file")
      continue
    # Directly update the submodule commit hash in the index.
    # See: https://stackoverflow.com/questions/33514642
    command = ["git", "update-index", "--cacheinfo", "160000", written, path]
    print("Updating", path, "to", written)
    utils.execute(command, cwd=repo_dir)


def init_submodules(repo_dir):
  print("*** Initializing submodules")
  utils.execute(["git", "submodule", "init"], cwd=repo_dir)


def parallel_shallow_update_submodules(repo_dir):
  print("*** Making shallow clone of submodules")
  utils.execute(["git", "submodule", "update", "--jobs", "8", "--depth", "1"],
                cwd=repo_dir)


def check_submodule_versions(repo_dir):
  diff_versions = get_diff_versions(repo_dir)
  if diff_versions:
    print("Submodule state differs from SUBMODULE_VERSIONS.txt file."
          " Run (and commit) one of:")
    print("  ./scripts/git/submodule_versions.py import"
          " # Use version in SUBMODULE_VERSIONS.txt ('written')")
    print("  ./scripts/git/submodule_versions.py export"
          " # Use version in git state ('actual')")
    for k, (current, written) in diff_versions.items():
      print(f"{k} : actual={current} written={written}")
    return False
  return True


def parse_arguments():
  parser = argparse.ArgumentParser()
  parser.add_argument("--repo", help="Repository root directory")
  parser.add_argument("command",
                      help="Command to run (show|import|export|check|init)")
  args = parser.parse_args()

  # Default repo path.
  if args.repo is None:
    args.repo = utils.find_git_toplevel()
  return args


def main(args):
  if args.command == "show":
    print(get_submodule_versions(args.repo))
  elif args.command == "export":
    sync_and_update_submodules(args.repo)
    export_versions(args.repo)
  elif args.command == "check":
    if not check_submodule_versions(args.repo):
      sys.exit(1)
  elif args.command == "import":
    import_versions(args.repo)
    sync_and_update_submodules(args.repo)
  elif args.command == "init":
    init_submodules(args.repo)
    # Redundant, since import_versions will only update if they differ,
    # but good to only print output about the import if it's actually
    # needed.
    if not check_submodule_versions(args.repo):
      print(
          "Warning: git submodule state does not match SUBMODULE_VERSIONS.txt. "
          "Using state in SUBMODULE_VERSIONS.txt")
      import_versions(args.repo)
    parallel_shallow_update_submodules(args.repo)
  else:
    print("Unrecognized command:", args.command)
    sys.exit(1)


if __name__ == "__main__":
  main(parse_arguments())