#!/usr/bin/env python3

import argparse
import os
import re
import subprocess
import sys
import tempfile

from typing import IO

BASE_DIR = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
sys.path.append(BASE_DIR)
from scripts.lib.zulip_tools import su_to_zulip, run

POSTGRES_USER = "postgres"

parser = argparse.ArgumentParser()
parser.add_argument("tarball", help="Filename of input tarball")


def restore_backup(tarball_file):
    # type: (IO[bytes]) -> None

    su_to_zulip(save_suid=True)

    import scripts.lib.setup_path_on_import

    # First, we unpack the /etc/zulip configuration, so we know how
    # this server is supposed to be configured (and can import
    # /etc/zulip/settings.py via `from django.conf import settings`,
    # next).  Ignore errors if zulip-backup/settings is not present
    # (E.g. because this is a development backup).
    tarball_file.seek(0, 0)
    subprocess.call(
        [
            "tar",
            "-C",
            "/etc/zulip",
            "--strip-components=2",
            "-xz",
            "zulip-backup/settings",
        ],
        stdin=tarball_file,
    )

    os.environ["DJANGO_SETTINGS_MODULE"] = "zproject.settings"
    from django.conf import settings

    paths = [
        ("settings", "/etc/zulip"),
        # zproject will only be present for development environment backups.
        ("zproject", os.path.join(settings.DEPLOY_ROOT, "zproject")),
    ]
    if settings.LOCAL_UPLOADS_DIR is not None:
        # We only need to restore LOCAL_UPLOADS_DIR if the system is
        # configured to locally host uploads.
        paths.append(("uploads", os.path.join(settings.DEPLOY_ROOT, settings.LOCAL_UPLOADS_DIR)))

    with tempfile.TemporaryDirectory(prefix="zulip-restore-backup-") as tmp:
        uid = os.getuid()
        gid = os.getgid()
        os.setresuid(0, 0, 0)
        for name, path in paths:
            os.makedirs(path, exist_ok=True)
            os.chown(path, uid, gid)
        os.setresuid(uid, uid, 0)

        assert not any("|" in name or "|" in path for name, path in paths)
        transform_args = [
            r"--transform=s|^zulip-backup/{}(/.*)?$|{}\1|x".format(
                re.escape(name), path.replace("\\", r"\\")
            )
            for name, path in paths
        ]

        os.mkdir(os.path.join(tmp, "zulip-backup"))
        tarball_file.seek(0, 0)
        run(["tar", "-C", tmp] + transform_args + ["-xPz"], stdin=tarball_file)

        # Now, extract the the database backup, destroy the old
        # database, and create a new, empty database.
        db_name = settings.DATABASES["default"]["NAME"]
        assert isinstance(db_name, str)
        db_dir = os.path.join(tmp, "zulip-backup", "database")
        os.setresuid(0, 0, 0)
        run(["chown", "-R", POSTGRES_USER, "--", tmp])
        run(
            [
                os.path.join(
                    settings.DEPLOY_ROOT, "scripts", "setup", "terminate-psql-sessions"
                ),
                "zulip",
                "zulip",
                "zulip_base",
            ]
        )
        as_postgres = ["su", "-s", "/usr/bin/env", "-", "--", POSTGRES_USER]
        run(as_postgres + ["dropdb", "--if-exists", "--", db_name])
        run(as_postgres + ["createdb", "-O", "zulip", "-T", "template0", "--", db_name])

        if settings.PRODUCTION:
            # In case we are restoring a backup from an older Zulip
            # version, there may be new secrets to generate.
            subprocess.check_call([
                os.path.join(settings.DEPLOY_ROOT, "scripts", "setup", "generate_secrets.py"),
                "--production",
            ])

            # If there is a local rabbitmq, we need to reconfigure it
            # to ensure the rabbitmq password matches the value in the
            # restored zulip-secrets.conf.  We need to be careful to
            # only do this if rabbitmq is configured to run locally on
            # the system.
            rabbitmq_host = subprocess.check_output(
                [os.path.join(settings.DEPLOY_ROOT,
                              "scripts", "get-django-setting"),
                 "RABBITMQ_HOST"]).strip().decode("utf-8")
            if rabbitmq_host in ["127.0.0.1", "::1", "localhost", "localhost6"]:
                run([os.path.join(settings.DEPLOY_ROOT,
                                  "scripts", "setup", "configure-rabbitmq")])

            # In production, we also need to do a `zulip-puppet-apply`
            # in order to apply any configuration from
            # /etc/zulip/zulip.conf to this system, since it was
            # originally installed without the restored copy of that
            # file.
            run(
                [
                    os.path.join(settings.DEPLOY_ROOT, "scripts", "zulip-puppet-apply"),
                    "-f",
                ]
            )

        # Now, restore the the database backup using pg_restore.  This
        # needs to run after zulip-puppet-apply to ensure full-text
        # search extensions are available and installed.
        run(as_postgres + ["pg_restore", "-d", db_name, "--", db_dir])
        run(["chown", "-R", str(uid), "--", tmp])
        os.setresuid(uid, uid, 0)

        if settings.PRODUCTION:
            run(["supervisorctl", "restart", "all"])

        run([os.path.join(settings.DEPLOY_ROOT, "scripts", "setup", "flush-memcached")])


if __name__ == "__main__":
    args = parser.parse_args()

    with open(args.tarball, "rb") as tarball_file:
        restore_backup(tarball_file)