# 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

import os
import re
import shlex
import shutil
import subprocess
from typing import Tuple

_repo_root = None


def get_repo_root() -> str:
    global _repo_root
    if _repo_root is None:
        _repo_root = os.getcwd()
        _validate_repo_root()
    return _repo_root


def get_submodule_root(submodule) -> str:
    path = os.path.join(get_repo_root(), "third_party", submodule)
    if not os.path.isdir(path):
        raise SystemExit(f"Could not find submodule: {path}")
    return path


def _validate_repo_root():
    # Look for something we know is there.
    known_dir = os.path.join(_repo_root, "compiler")
    if not os.path.isdir(known_dir):
        raise SystemExit(f"ERROR: Must run from the iree repository root. "
                         f"Actually in: {_repo_root}")


def git_setup_remote(remote_alias, url, *, repo_dir=None):
    needs_create = False
    try:
        existing_url = git_exec(["remote", "get-url", remote_alias],
                                capture_output=True,
                                repo_dir=repo_dir,
                                quiet=True)
        existing_url = existing_url.strip()
        if existing_url == url:
            return
    except subprocess.CalledProcessError:
        # Does not exist.
        needs_create = True

    if needs_create:
        git_exec(["remote", "add", "--no-tags", remote_alias, url],
                 repo_dir=repo_dir)
    else:
        git_exec(["remote", "set-url", remote_alias, url], repo_dir=repo_dir)


def git_is_porcelain(*, repo_dir=None):
    output = git_exec(["status", "--porcelain", "--untracked-files=no"],
                      capture_output=True,
                      quiet=True,
                      repo_dir=repo_dir).strip()
    return not bool(output)


def git_check_porcelain(*, repo_dir=None):
    output = git_exec(["status", "--porcelain", "--untracked-files=no"],
                      capture_output=True,
                      quiet=True,
                      repo_dir=repo_dir).strip()
    if output:
        actual_repo_dir = get_repo_root() if repo_dir is None else repo_dir
        raise SystemExit(
            f"ERROR: git directory {actual_repo_dir} is not clean. "
            f"Please stash changes:\n{output}")


def git_fetch(*, repository=None, ref=None, repo_dir=None):
    args = ["fetch"]
    if repository:
        args.append(repository)
    if ref is not None:
        args.append(ref)
    git_exec(args, repo_dir=repo_dir)


def git_checkout(ref, *, repo_dir=None):
    git_exec(["checkout", ref], repo_dir=repo_dir)


def git_create_branch(branch_name,
                      *,
                      checkout=True,
                      ref=None,
                      force=False,
                      repo_dir=None):
    branch_args = ["branch"]
    if force:
        branch_args.append("-f")
    branch_args.append(branch_name)
    if ref is not None:
        branch_args.append(ref)
    git_exec(branch_args, repo_dir=repo_dir)

    if checkout:
        git_exec(["checkout", branch_name], repo_dir=repo_dir)


def git_push_branch(repository, branch_name, *, force=False, repo_dir=None):
    push_args = ["push", "--set-upstream"]
    if force:
        push_args.append("-f")
    push_args.append(repository)
    push_args.append(f"{branch_name}:{branch_name}")
    git_exec(push_args, repo_dir=repo_dir)


def git_branch_exists(branch_name, *, repo_dir=None):
    output = git_exec(["branch", "-l", branch_name],
                      repo_dir=repo_dir,
                      quiet=True,
                      capture_output=True).strip()
    return bool(output)


def git_submodule_set_origin(path, *, url=None, branch=None, repo_dir=None):
    if url is not None:
        git_exec(["submodule", "set-url", "--", path, url], repo_dir=repo_dir)

    if branch is not None:
        try:
            if branch == "--default":
                git_exec(["submodule", "set-branch", "--default", "--", path],
                         repo_dir=repo_dir)
            else:
                git_exec([
                    "submodule", "set-branch", "--branch", branch, "--", path
                ],
                         repo_dir=repo_dir)
        except subprocess.CalledProcessError:
            # The set-branch command returns 0 on change and !0 on no change.
            # This is a bit unfortunate.
            ...


def git_reset(ref, *, hard=True, repo_dir=None):
    args = ["reset"]
    if hard:
        args.append("--hard")
    args.append(ref)
    git_exec(args, repo_dir=repo_dir)


def git_current_commit(*, repo_dir=None) -> Tuple[str, str]:
    output = git_exec(["log", "-n", "1", "--pretty=format:%H (%ci): %s"],
                      capture_output=True,
                      repo_dir=repo_dir,
                      quiet=True)
    output = output.strip()
    parts = output.split(" ")
    # Return commit, full_summary
    return parts[0], output


def git_create_commit(*, message, add_all=False, repo_dir=None):
    if add_all:
        git_exec(["add", "-A"], repo_dir=repo_dir)
    git_exec(["commit", "-m", message])


def git_ls_remote_branches(repository_url, *, filter=None, repo_dir=None):
    args = ["ls-remote", "-h", repository_url]
    if filter:
        args.extend(filter)
    output = git_exec(args, quiet=True, capture_output=True)
    lines = output.strip().splitlines(keepends=False)

    # Format is <commit> refs/heads/branch_name
    def extract_branch(line):
        parts = re.split("\\s+", line)
        ref = parts[1]
        prefix = "refs/heads/"
        if ref.startswith(prefix):
            ref = ref[len(prefix):]
        return ref

    return [extract_branch(l) for l in lines]


def git_exec(args, *, repo_dir=None, quiet=False, capture_output=False):
    full_args = ["git"] + args
    full_args_quoted = [shlex.quote(a) for a in full_args]
    if not repo_dir:
        repo_dir = get_repo_root()
    if not quiet:
        print(f"  ++ EXEC: (cd {repo_dir} && {' '.join(full_args_quoted)})")
    if capture_output:
        return subprocess.check_output(full_args, cwd=repo_dir).decode("utf-8")
    else:
        subprocess.check_call(full_args, cwd=repo_dir)