#!/usr/bin/python3

"""
The UBports QA scripts allow you to efficiently manage PPAs from repo.ubports.com for testing deb components.
Copyright 2018 UBports Foundation
Licensed GPL v. 3 or later
"""
import subprocess
import os
import argparse
from enum import Enum
import requests
import logging


NO_PR = -1
GITHUB_API_PULLREQUEST = "https://api.github.com/repos/ubports/{repo}/pulls/{num}"
JENKINS_API_BUILD = "https://ci.ubports.com/blue/rest/organizations/jenkins/pipelines/{repo}/branches/{ref}"

NO_UPGRADE_PREF = """\
Package: *
Pin: origin repo.ubports.com
Pin: release o=UBports,a=focal
Pin-Priority: 50

Package: *
Pin: release o=Ubuntu
Pin-Priority: 10
"""
NO_UPGRADE_PREF_PATH = "/etc/apt/preferences.d/aaa-ubports-no-upgrade.pref"

LOG = logging.getLogger(name="ubports-qa")
IS_ROOT = os.geteuid() == 0


class Status(Enum):
    SUCCESS = 1
    BUILDING = 2
    FAILED = 3


class WritableRootFS:
    """
    A class to be used with the `with` statement to mount `/` read-write, for example::

        with WritableRootFS():
            write_file(/good_file)

    `/` will be remounted read-only on close, unless the file /userdata/.writable_image
    exists.
    """

    def __enter__(self):
        self.attempt_writable_mount()

    def __exit__(self, exc_type, value, traceback):
        self.attempt_unmount()

    @classmethod
    def attempt_writable_mount(cls):
        """Tries to mount the rootfs read-write"""
        LOG.debug("Attempting to mount rootfs read-write.")
        ensure_root()
        subprocess.run(["mount", "-o", "rw,remount", "/"])
        LOG.debug("Making sure we have enough free space for archives.")
        # Make sure this directory exists. Sometimes it doesn't on a freshly-flashed system.
        os.makedirs("/var/cache/apt/archives", exist_ok=True)
        subprocess.run(["mount", "-t", "tmpfs", "tmpfs", "/var/cache/apt/archives"])
        LOG.debug("Making sure we have enough free space for the package lists.")
        # Make sure this directory exists. Sometimes it doesn't on a freshly-flashed system.
        os.makedirs("/var/lib/apt/lists", exist_ok=True)
        subprocess.run(["mount", "-t", "tmpfs", "tmpfs", "/var/lib/apt/lists"])
        

    @classmethod
    def attempt_unmount(cls):
        LOG.debug("Unmounting tmpfs")
        ensure_root()
        subprocess.run(["umount", "/var/cache/apt/archives"])
        subprocess.run(["umount", "/var/lib/apt/lists"])
        LOG.debug("Attempting to mount rootfs read-only.")
        if not os.path.exists("/userdata/.writable_image"):
            os.sync()
            try:
                subprocess.run(["mount", "-o", "ro,remount", "/"], check=True)
            except subprocess.CalledProcessError:
                LOG.warning("Failed to remount root filesystem read-only.")
                LOG.warning("Please consider rebooting your device.")


def ensure_root():
    if not IS_ROOT:
        die("Insufficient permissions, please run with sudo.")


def apt_update():
    LOG.debug("Running 'apt update'.")
    try:
        subprocess.run(["apt", "update"], check=True)
    except subprocess.CalledProcessError:
        LOG.error("Failed to run 'apt update'. See the output above for details.")


def apt_upgrade():
    LOG.debug("Running 'apt full-upgrade'.")
    try:
        subprocess.run(["apt", "full-upgrade"], check=True)
        subprocess.run(["apt", "autoremove"], check=True)
    except subprocess.CalledProcessError:
        LOG.error("Failed to run 'apt full-upgrade'. See the output above for details.")


def get_list_file(branch):
    list_file = "/etc/apt/sources.list.d/ubports-{}.list".format(branch)
    LOG.debug("List file is " + list_file)
    return list_file


def get_pref_file(branch):
    pref_file = "/etc/apt/preferences.d/ubports-{}.pref".format(branch)
    LOG.debug("Pref file is " + "/etc/apt/preferences.d/ubports-{}.pref".format(branch))
    return pref_file


def list_exists(branch):
    return os.path.isfile(get_list_file(branch))


def pref_exists(branch):
    return os.path.isfile(get_pref_file(branch))


def list_lists():
    return [
        f.split("ubports-")[1].split(".list")[0]
        for f in os.listdir("/etc/apt/sources.list.d/")
        if os.path.isfile("/etc/apt/sources.list.d/" + f) and f.startswith("ubports-")
    ]


def add_list(branch):
    if list_exists(branch):
        return
    base = "http://repo.ubports.com/"
    if requests.head("{}dists/{}/".format(base, branch) ).status_code != 200:
        die("PPA not found")
    with open(get_list_file(branch), "w+") as repo_list:
        repo_list.write("deb {} {} main".format(base, branch))


def add_pref(branch):
    if pref_exists(branch):
        return
    with open(get_pref_file(branch), "w+") as pref:
        # TODO: edit and set own pin priority to allow up/down/side-grade without
        # removal of repo etc
        pref.write("Package: *\n"
                   "Pin: release o=UBports,a={}\n"
                   "Pin-Priority: 3001\n".format(branch))

def add_no_upgrade_pref():
    with open(NO_UPGRADE_PREF_PATH, "w+") as pref:
        pref.write(NO_UPGRADE_PREF)

def remove_no_upgrade_pref():
    os.remove(NO_UPGRADE_PREF_PATH)

def remove_list(branch):
    # If it does not exist, just ignore
    if not list_exists(branch):
        return
    os.remove(get_list_file(branch))


def remove_pref(branch):
    # If it does not exist, just ignore
    if not pref_exists(branch):
        return
    os.remove(get_pref_file(branch))


def get_github_pr(repo, num):
    pull_request_url = GITHUB_API_PULLREQUEST.format(repo=repo, num=num)
    LOG.debug("Getting pull request information from " + pull_request_url)
    ret = requests.get(pull_request_url)
    if ret.status_code != 200:
        die("Pull-Request not found")
    return ret.json()


def get_jenkins_build(repo, ref):
    LOG.debug("Checking PR for valid build on UBports CI.")
    LOG.debug(JENKINS_API_BUILD.format(repo=repo, ref=ref))
    ret = requests.get(JENKINS_API_BUILD.format(repo=repo, ref=ref))
    if ret.status_code != 200:
        die("Jenkins build not found")
    return ret.json()


def get_issue_status(repo, ref):
    # Prepend "ubports/" because all CI jobs start with it.
    build = get_jenkins_build("ubports/" + repo, ref)["latestRun"]
    LOG.debug(build)
    if build["result"] == "SUCCESS":
        return Status.SUCCESS
    elif build["result"] == "BULDING":
        return Status.BULDING

    return Status.FAILED


def die(error_message):
    """Logs error_message and exits with status 3"""
    LOG.critical(error_message)
    exit(3)


def install_command(args):
    """Install a PPA or Pull Request"""
    if args.pr is not NO_PR:
        # Support both with or without ubports/ prefix.
        github_repo = args.repo.replace("ubports/", "")
        ref = "PR-" + str(args.pr)
        LOG.debug(
            "Checking repository ubports/" + github_repo + " for PR #" + str(args.pr)
        )
        status = get_issue_status(github_repo, ref)
        if status == Status.FAILED:
            die("Issue failed to build")
        if status == Status.BUILDING:
            die("Issue is currently building")
        repository_name = "PR_{repo}_{pr_num}".format(repo=github_repo, pr_num=args.pr)
    else:
        repository_name = args.repo
        LOG.debug("No PR set, installing " + repository_name)

    with WritableRootFS():
        add_list(repository_name)
        add_pref(repository_name)
        try:
            add_no_upgrade_pref()
            apt_update()
            apt_upgrade()
        finally:
            remove_no_upgrade_pref()

    LOG.info(
        "You can remove this repository by running 'sudo ubports-qa remove {}'".format(
            repository_name
        )
    )


def remove_command(args):
    """Remove and uninstall a PPA"""
    repository = args.repo
    LOG.info("Removing repo " + repository)
    if not list_exists(repository):
        die("Repo {} is not installed".format(repository))
    with WritableRootFS():
        remove_list(repository)
        remove_pref(repository)
        apt_update()
        apt_upgrade()


def list_command(args):
    """List installed PPAs"""
    print("\n".join(list_lists()))


def update_command(args):
    """Update all packages using apt"""
    with WritableRootFS():
        apt_update()
        apt_upgrade()


parser = argparse.ArgumentParser(
    description="The UBports QA scripts allow you to efficiently manage PPAs from repo.ubports.com for testing deb components. See http://docs.ubports.com/en/latest/about/process/ppa.html."
)
parser.add_argument(
    "-v", "--verbose", help="Print verbose logging information", action="store_true"
)
subparsers = parser.add_subparsers(help="")

parser_install = subparsers.add_parser(
    "install",
    help=install_command.__doc__,
    description="Install a ppa or pull-request. See http://docs.ubports.com/en/latest/about/process/ppa.html.",
)
parser_install.add_argument(
    "repo",
    type=str,
    help="Name of a PPA on repo.ubports.com. Note that for certain branches on github, PPAs are automatically created. For example branch names that start with 'xenial_-_' have such automatic PPAs. Example: 'ubports-qa install xenial_-_sessionrestore'. Alternatively, see below how 'repo' is interpreted if the 'pr' argument is used."
)
parser_install.add_argument(
    "pr",
    type=int,
    help="Numeric ID of a pull-request. When this 'pr' parameter is used, then the 'repo' parameter is interpreted as the name of a github repository. Example: 'ubports-qa calendar-app 123'",
    nargs="?",
    default=NO_PR,
)
parser_install.set_defaults(func=install_command)

parser_remove = subparsers.add_parser(
    "remove",
    aliases=["uninstall"],
    help=remove_command.__doc__,
    description="Remove and uninstall a ppa",
)
parser_remove.add_argument("repo", type=str, help="Name of the ppa")
parser_remove.set_defaults(func=remove_command)

parser_list = subparsers.add_parser(
    "list", help=list_command.__doc__, description="List installed PPAs"
)
parser_list.set_defaults(func=list_command)

parser_update = subparsers.add_parser(
    "update", help=update_command.__doc__, description="Update all packages using apt"
)
parser_update.set_defaults(func=update_command)

try:
    ARGS = parser.parse_args()
    logging.basicConfig()
    if ARGS.verbose:
        LOG.setLevel(logging.DEBUG)
    else:
        LOG.setLevel(logging.INFO)
    ARGS.func(ARGS)
except IOError as e:
    # We weren't allowed to do something with a file.
    # Either we aren't root or the disk is read-only.
    ensure_root()
    die(e)
except AttributeError as e:
    # The user typed an incorrect command
    parser.print_help()
    exit(4)
