#!/usr/bin/env python3
import argparse
import os
import subprocess
import sys
import urllib
from collections import namedtuple
from fnmatch import fnmatch

import requests

try:
    import semver
except ImportError:
    print("Missing required library: semver.")
    exit(1)

REDASH_HOME = os.environ.get("REDASH_HOME", "/opt/redash")
CURRENT_VERSION_PATH = "{}/current".format(REDASH_HOME)


def run(cmd, cwd=None):
    if not cwd:
        cwd = REDASH_HOME

    return subprocess.check_output(cmd, cwd=cwd, shell=True, stderr=subprocess.STDOUT)


def confirm(question):
    reply = str(input(question + " (y/n): ")).lower().strip()

    if reply[0] == "y":
        return True
    if reply[0] == "n":
        return False
    else:
        return confirm("Please use 'y' or 'n'")


def version_path(version_name):
    return "{}/{}".format(REDASH_HOME, version_name)


END_CODE = "\033[0m"


def colored_string(text, color):
    if sys.stdout.isatty():
        return "{}{}{}".format(color, text, END_CODE)
    else:
        return text


def h1(text):
    print(colored_string(text, "\033[4m\033[1m"))


def green(text):
    print(colored_string(text, "\033[92m"))


def red(text):
    print(colored_string(text, "\033[91m"))


class Release(namedtuple("Release", ("version", "download_url", "filename", "description"))):
    def v1_or_newer(self):
        return semver.compare(self.version, "1.0.0-alpha") >= 0

    def is_newer(self, version):
        return semver.compare(self.version, version) > 0

    @property
    def version_name(self):
        return self.filename.replace(".tar.gz", "")


def get_latest_release_from_ci():
    response = requests.get(
        "https://circleci.com/api/v1.1/project/github/getredash/redash/latest/artifacts?branch=master"
    )

    if response.status_code != 200:
        exit("Failed getting releases (status code: %s)." % response.status_code)

    tarball_asset = filter(lambda asset: asset["url"].endswith(".tar.gz"), response.json())[0]
    filename = urllib.unquote(tarball_asset["pretty_path"].split("/")[-1])
    version = filename.replace("redash.", "").replace(".tar.gz", "")

    release = Release(version, tarball_asset["url"], filename, "")

    return release


def get_release(channel):
    if channel == "ci":
        return get_latest_release_from_ci()

    response = requests.get("https://version.redash.io/api/releases?channel={}".format(channel))
    release = response.json()[0]

    filename = release["download_url"].split("/")[-1]
    release = Release(release["version"], release["download_url"], filename, release["description"])

    return release


def link_to_current(version_name):
    green("Linking to current version...")
    run("ln -nfs {} {}".format(version_path(version_name), CURRENT_VERSION_PATH))


def restart_services():
    # We're doing this instead of simple 'supervisorctl restart all' because
    # otherwise it won't notice that /opt/redash/current pointing at a different
    # directory.
    green("Restarting...")
    try:
        run("sudo /etc/init.d/redash_supervisord restart")
    except subprocess.CalledProcessError as e:
        run("sudo service supervisor restart")


def update_requirements(version_name):
    green("Installing new Python packages (if needed)...")
    new_requirements_file = "{}/requirements.txt".format(version_path(version_name))

    install_requirements = False

    try:
        run("diff {}/requirements.txt {}".format(CURRENT_VERSION_PATH, new_requirements_file)) != 0
    except subprocess.CalledProcessError as e:
        if e.returncode != 0:
            install_requirements = True

    if install_requirements:
        run("sudo pip install -r {}".format(new_requirements_file))


def apply_migrations(release):
    green("Running migrations (if needed)...")
    if not release.v1_or_newer():
        return apply_migrations_pre_v1(release.version_name)

    run("sudo -u redash bin/run ./manage.py db upgrade", cwd=version_path(release.version_name))


def find_migrations(version_name):
    current_migrations = set(
        [f for f in os.listdir("{}/migrations".format(CURRENT_VERSION_PATH)) if fnmatch(f, "*_*.py")]
    )
    new_migrations = sorted(
        [f for f in os.listdir("{}/migrations".format(version_path(version_name))) if fnmatch(f, "*_*.py")]
    )

    return [m for m in new_migrations if m not in current_migrations]


def apply_migrations_pre_v1(version_name):
    new_migrations = find_migrations(version_name)

    if new_migrations:
        green("New migrations to run: ")
        print(", ".join(new_migrations))
    else:
        print("No new migrations in this version.")

    if new_migrations and confirm("Apply new migrations? (make sure you have backup)"):
        for migration in new_migrations:
            print("Applying {}...".format(migration))
            run(
                "sudo sudo -u redash PYTHONPATH=. bin/run python migrations/{}".format(migration),
                cwd=version_path(version_name),
            )


def download_and_unpack(release):
    directory_name = release.version_name

    green("Downloading release tarball...")
    run(
        'sudo wget --header="Accept: application/octet-stream" -O {} {}'.format(release.filename, release.download_url)
    )
    green("Unpacking to: {}...".format(directory_name))
    run("sudo mkdir -p {}".format(directory_name))
    run("sudo tar -C {} -xvf {}".format(directory_name, release.filename))

    green("Changing ownership to redash...")
    run("sudo chown redash {}".format(directory_name))

    green("Linking .env file...")
    run("sudo ln -nfs {}/.env {}/.env".format(REDASH_HOME, version_path(directory_name)))


def current_version():
    real_current_path = os.path.realpath(CURRENT_VERSION_PATH).replace(".b", "+b")
    return real_current_path.replace(REDASH_HOME + "/", "").replace("redash.", "")


def verify_minimum_version():
    green("Current version: " + current_version())
    if semver.compare(current_version(), "0.12.0") < 0:
        red("You need to have Redash v0.12.0 or newer to upgrade to post v1.0.0 releases.")
        green("To upgrade to v0.12.0, run the upgrade script set to the legacy channel (--channel legacy).")
        exit(1)


def show_description_and_confirm(description):
    if description:
        print(description)

        if not confirm("Continue with upgrade?"):
            red("Cancelling upgrade.")
            exit(1)


def verify_newer_version(release):
    if not release.is_newer(current_version()):
        red("The found release is not newer than your current deployed release ({}).".format(current_version()))
        if not confirm("Continue with upgrade?"):
            red("Cancelling upgrade.")
            exit(1)


def deploy_release(channel):
    h1("Starting Redash upgrade:")

    release = get_release(channel)
    green("Found version: {}".format(release.version))

    if release.v1_or_newer():
        verify_minimum_version()

    verify_newer_version(release)
    show_description_and_confirm(release.description)

    try:
        download_and_unpack(release)
        update_requirements(release.version_name)
        apply_migrations(release)
        link_to_current(release.version_name)
        restart_services()
        green("Done! Enjoy.")
    except subprocess.CalledProcessError as e:
        red("Failed running: {}".format(e.cmd))
        red("Exit status: {}\nOutput:\n{}".format(e.returncode, e.output))


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--channel", help="The channel to get release from (default: stable).", default="stable")
    args = parser.parse_args()

    deploy_release(args.channel)
