#!/usr/bin/env python3
import argparse
import logging
import os
import subprocess
import sys
import time

LOCAL_GIT_CACHE_DIR = "/srv/zulip.git"
os.environ["PYTHONUNBUFFERED"] = "y"

sys.path.append(os.path.join(os.path.dirname(__file__), "..", ".."))
from scripts.lib.zulip_tools import (
    DEPLOYMENTS_DIR,
    assert_running_as_root,
    get_config,
    get_config_file,
    get_deploy_options,
    get_deploy_root,
    get_deployment_lock,
    make_deploy_path,
    overwrite_symlink,
    release_deployment_lock,
    su_to_zulip,
)

config_file = get_config_file()
deploy_options = get_deploy_options(config_file)
remote_url = get_config(
    config_file, "deployment", "git_repo_url", "https://github.com/zulip/zulip.git"
)

assert_running_as_root(strip_lib_from_paths=True)

# make sure we have appropriate file permissions
os.umask(0o22)

logging.Formatter.converter = time.gmtime
logging.basicConfig(format="%(asctime)s upgrade-zulip-from-git: %(message)s", level=logging.INFO)

parser = argparse.ArgumentParser()
parser.add_argument("refname", help="Git reference, e.g. a branch, tag, or commit ID.")
git_ref = parser.add_mutually_exclusive_group()
git_ref.add_argument(
    "--remote-url", help="Override the Git remote URL configured in /etc/zulip/zulip.conf."
)
git_ref.add_argument(
    "--local-ref",
    action="store_true",
    help="Provided branch name has been pushed directly to /srv/zulip.git already",
)
args, extra_options = parser.parse_known_args()

refname = args.refname
# Command line remote URL will be given preference above the one
# in /etc/zulip/zulip.conf.
if args.remote_url:
    remote_url = args.remote_url

os.makedirs(DEPLOYMENTS_DIR, exist_ok=True)

error_rerun_script = f"{DEPLOYMENTS_DIR}/current/scripts/upgrade-zulip-from-git {refname}"
get_deployment_lock(error_rerun_script)

try:
    deploy_path = make_deploy_path()

    # Populate LOCAL_GIT_CACHE_DIR with both the requested remote and zulip/zulip.
    if not os.path.exists(LOCAL_GIT_CACHE_DIR):
        logging.info("Making local repository cache")
        subprocess.check_call(
            ["git", "init", "--bare", "-q", LOCAL_GIT_CACHE_DIR],
            stdout=subprocess.DEVNULL,
        )
        subprocess.check_call(
            ["git", "remote", "add", "origin", remote_url],
            cwd=LOCAL_GIT_CACHE_DIR,
        )

    if os.stat(LOCAL_GIT_CACHE_DIR).st_uid == 0:
        subprocess.check_call(["chown", "-R", "zulip:zulip", LOCAL_GIT_CACHE_DIR])

    os.chdir(LOCAL_GIT_CACHE_DIR)
    if not args.local_ref:
        subprocess.check_call(
            ["git", "remote", "set-url", "origin", remote_url], preexec_fn=su_to_zulip
        )

    fetch_spec = subprocess.check_output(
        ["git", "config", "remote.origin.fetch"],
        preexec_fn=su_to_zulip,
        text=True,
    ).strip()
    if fetch_spec in ("+refs/*:refs/*", "+refs/heads/*:refs/heads/*"):
        # The refspec impinges on refs/heads/ -- this is an old mirror
        # configuration.
        logging.info("Cleaning up mirrored repository")
        # remotes.origin.mirror may not be set -- we do not use
        # check_call to ignore errors if it's already missing
        subprocess.call(
            ["git", "config", "--unset", "remote.origin.mirror"],
            preexec_fn=su_to_zulip,
        )
        subprocess.check_call(
            ["git", "config", "remote.origin.fetch", "+refs/heads/*:refs/remotes/origin/*"],
            preexec_fn=su_to_zulip,
        )
        matching_refs = subprocess.check_output(
            ["git", "for-each-ref", "--format=%(refname)", "refs/pull/", "refs/heads/"],
            preexec_fn=su_to_zulip,
            text=True,
        ).splitlines()

        # We can't use `git worktree list --porcelain -z` here because
        # Ubuntu 22.04 only has git 2.34.1, and -z was
        # introduced in 2.36
        worktree_data = subprocess.check_output(
            ["git", "worktree", "list", "--porcelain"],
            preexec_fn=su_to_zulip,
            text=True,
        ).splitlines()
        keep_refs = set()
        for worktree_line in worktree_data:
            if worktree_line.startswith("branch "):
                keep_refs.add(worktree_line.removeprefix("branch "))

        delete_input = "".join(
            f"delete {refname}\n" for refname in matching_refs if refname not in keep_refs
        )
        subprocess.run(
            ["git", "update-ref", "--stdin"],
            check=True,
            preexec_fn=su_to_zulip,
            input=delete_input,
            text=True,
        )

        logging.info("Repacking repository after pruning unnecessary refs...")
        subprocess.check_call(
            ["git", "gc", "--prune=now"],
            preexec_fn=su_to_zulip,
        )

    if not args.local_ref:
        subprocess.check_call(
            [os.path.join(get_deploy_root(), "scripts/lib/update-git-upstream")],
            preexec_fn=su_to_zulip,
        )

    # Generate the deployment directory via git worktree from our local repository.
    try:
        fullref = f"refs/tags/{refname}"
        commit_hash = subprocess.check_output(
            ["git", "rev-parse", "--verify", fullref],
            preexec_fn=su_to_zulip,
            text=True,
            stderr=subprocess.DEVNULL,
        ).strip()
    except subprocess.CalledProcessError as e:
        if e.returncode == 128:
            # Try in the origin namespace, or local heads if --local-ref
            if args.local_ref:
                fullref = f"refs/heads/{refname}"
            else:
                fullref = f"refs/remotes/origin/{refname}"
            commit_hash = subprocess.check_output(
                ["git", "rev-parse", "--verify", fullref],
                preexec_fn=su_to_zulip,
                text=True,
                stderr=subprocess.DEVNULL,
            ).strip()
    refname = fullref
    subprocess.check_call(
        ["git", "worktree", "add", "--detach", deploy_path, refname],
        stdout=subprocess.DEVNULL,
        preexec_fn=su_to_zulip,
    )
    os.chdir(deploy_path)
    extra_flags = []
    if not refname.startswith("refs/tags/"):
        extra_flags = ["-t"]
    subprocess.check_call(
        [
            "git",
            "checkout",
            "-q",
            *extra_flags,
            "-b",
            "deployment-" + os.path.basename(deploy_path),
            refname,
        ],
        preexec_fn=su_to_zulip,
    )

    overwrite_symlink("/etc/zulip/settings.py", "zproject/prod_settings.py")

    overwrite_symlink(deploy_path, os.path.join(DEPLOYMENTS_DIR, "next"))

    try:
        subprocess.check_call(
            [
                os.path.join(deploy_path, "scripts", "lib", "upgrade-zulip-stage-2"),
                deploy_path,
                "--from-git",
                *deploy_options,
                *extra_options,
            ]
        )
    except subprocess.CalledProcessError:
        # There's no use in showing a stacktrace here; it just hides
        # the error from stage 2.
        sys.exit(1)
finally:
    release_deployment_lock()
