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

TARBALL_ARCHIVE_PATH = "/home/zulip/archives"
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_file,
    get_deploy_options,
    get_deployment_lock,
    release_deployment_lock,
    su_to_zulip,
)

config_file: configparser.RawConfigParser = get_config_file()
deploy_options = get_deploy_options(config_file)

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: %(message)s", level=logging.INFO)

parser = argparse.ArgumentParser()
parser.add_argument("tarball", help="Path to Zulip Server tarball")
args, extra_options = parser.parse_known_args()

error_rerun_script = f"{DEPLOYMENTS_DIR}/current/scripts/upgrade-zulip {args.tarball}"
get_deployment_lock(error_rerun_script)

try:
    # Copy the release tarball to an archival path that's readable by
    # the Zulip user, and then unpack it from that directory, so that
    # we can unpack using the Zulip user even if the original path was
    # not readable by the Zulip user.
    logging.info("Archiving the tarball under %s", TARBALL_ARCHIVE_PATH)
    os.makedirs(TARBALL_ARCHIVE_PATH, exist_ok=True)
    archived_tarball_path = os.path.join(TARBALL_ARCHIVE_PATH, os.path.basename(args.tarball))
    shutil.copy(args.tarball, archived_tarball_path)
    subprocess.check_output(["chown", "-R", "zulip:zulip", TARBALL_ARCHIVE_PATH])

    logging.info("Unpacking the tarball")
    unpack_zulip = os.path.realpath(os.path.join(os.path.dirname(__file__), "unpack-zulip"))
    deploy_path = subprocess.check_output(
        [unpack_zulip, archived_tarball_path], preexec_fn=su_to_zulip, text=True
    )

    # Chdir to deploy_path and then run upgrade-zulip-stage-2 from the
    # new version of Zulip (having the upgrade logic run from the new
    # version is much better for fixing bugs in the upgrade process).
    deploy_path = deploy_path.strip()
    os.chdir(deploy_path)
    try:
        subprocess.check_call(
            [
                os.path.abspath("./scripts/lib/upgrade-zulip-stage-2"),
                deploy_path,
                *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()