#!/usr/bin/python3

import argparse
import os
import subprocess
import sys
import tempfile

def main():
    parser = argparse.ArgumentParser(
        description="pushpkg, push aosc package to repo.aosc.io or mirrors"
    )
    parser.add_argument(
        "username", metavar="USERNAME", type=str, help="Your LDAP username.", nargs="?"
    )
    parser.add_argument(
        "branch",
        metavar="BRANCH",
        type=str,
        help="AOSC OS update branch (stable, stable-proposed, testing, etc.)",
        nargs="?",
    )
    parser.add_argument(
        "component",
        metavar="COMPONENT",
        type=str,
        help="(Optional) Repository component (main, bsp-sunxi, etc.) "
        'Falls back to "main" if not specified.',
        nargs="?",
        default="main",
    )
    parser.add_argument(
        "-v",
        "--verbose",
        action="store_true",
        help="Enable verbose logging for ssh and rsync",
    )
    parser.add_argument(
        "-d",
        "--delete",
        action="store_true",
        help="Clean OUTPUT directory after finishing uploading.",
    )
    parser.add_argument(
        "-f",
        "--force-push-noarch-package",
        action="store_true",
        help="Force Push noarch package.",
    )
    parser.add_argument(
        "-r", "--retro", action="store_true", help="Push to AOSC OS/Retro repo"
    )
    parser.add_argument(
        "-6", "--ipv6", action="store_true", help="Use IPv6 addresses only"
    )
    parser.add_argument(
        "-4", "--ipv4", action="store_true", help="Use IPv4 addresses only"
    )
    parser.add_argument(
        "-C", "--compress", action="store_true", help="Requests compression of all data for SSH, note this option should be used with caution"
    )
    parser.add_argument(
        "--host", type=str, nargs='?', default='repo.aosc.io', help="Specify the rsync host to push packages, defaults to repo.aosc.io"
    )
    parser.add_argument(
        "-i", "--identity-file", type=str, help="SSH identity file", default=None
    )
    parser.add_argument(
        "--allow-marking-upload-done-to-fail", action="store_true", help="Allow marking upload as done to fail"
    )
    args = parser.parse_args()
    username = args.username
    branch = args.branch
    component = args.component
    if not username:
        username = detect_and_ask(
            "username",
            os.environ.get("PUSHPKG_USERNAME")
            or subprocess.check_output(["whoami"], shell=True).strip().decode(),
        )
    if not branch:
        branch = detect_and_ask(
            "branch", os.path.split(os.getcwd())[1].lstrip("OUTPUT-")
        )
    if not os.path.isdir("./debs"):
        print("[!!!] debs is not a directory!")
        sys.exit(1)
    delete_junk()
    partial_dir = f"/mirror/.debs-incoming/pushpkg.{branch}_{next(tempfile._get_candidate_names())}"
    upload_url = make_upload_url(username, branch, component, args.retro, args.host)
    rsync_non_noarch_file(upload_url, branch, partial_dir, args.identity_file, args.verbose, args.ipv6, args.ipv4, args.compress)
    if have_noarch_files():
        rsync_noarch_file(upload_url, branch, partial_dir, args.identity_file, args.verbose, args.force_push_noarch_package, args.ipv6, args.ipv4, args.compress)
    else:
        print("[+] There is no noarch packages. Skipping.")
    if args.delete:
        clean_output_directory()
    mark_upload_done(username, partial_dir, args.host, args.identity_file, args.verbose, args.allow_marking_upload_done_to_fail, args.ipv6, args.ipv4, args.compress)


def detect_and_ask(type_name: str, arg: str) -> str:
    choice = input(f"[+] Detected {type_name} {arg}. Use this one? [Y/n] ").lower()

    if choice == "n":
        return input(f"[+] Please fill in the {type_name}: ")
    if choice in ("y", ""):
        return arg

    print("[!!!] Unexpected response!")
    return detect_and_ask(type_name, arg)


def make_upload_url(username: str, branch: str, component: str, is_retro: bool, host: str) -> str:
    return f"{username}@{host}:/mirror/debs{'-retro' if is_retro else ''}/pool/{branch}/{component}/"


def delete_junk():
    print("[+] Removing loose files ...")
    debs_path = os.path.abspath("./debs")
    command = [
        "sudo",
        "find",
        debs_path,
        "-maxdepth",
        "1",
        "-type",
        "f",
        "-delete",
        "-print",
    ]
    subprocess.check_call(command)


def mark_upload_done(username: str, partial_dir: str, host: str, identity_file=None, verbose=False, allow_fail=False, ipv6=False, ipv4=False, compress=False):
    command = ["ssh", f"{username}@{host}", "touch", "/mirror/.updated", "&&" ,"rm" , "-rf", partial_dir]
    if verbose:
        command.insert(1, "-v")
    if identity_file:
        command.insert(1, "-i")
        command.insert(2, identity_file)
    if ipv4:
        command.insert(1, "-4")
    if ipv6:
        command.insert(1, "-6")
    if compress:
        command.insert(1, "-C")
    try:
        subprocess.check_call(command)
    except subprocess.CalledProcessError as err:
        if allow_fail:
            print('Marking upload done failed with "', err, '", ignored')
        else:
            raise err


def rsync_non_noarch_file(upload_url: str, branch: str, partial_dir: str, identity_file=None, verbose=False, ipv6=False, ipv4=False, compress=False):
    print("[+] Uploading arch-specific packages ...")
    flags = []
    ssh_option = ""

    if compress:
        ssh_option = "-C"

    if identity_file:
        flags.extend(["--rsh", f"ssh {ssh_option} -i {identity_file}"])

    if branch == "stable":
        flags.append("--ignore-existing")

    if verbose:
        flags.append("--verbose")

    if ipv4:
        flags.append("--ipv4")
    if ipv6:
        flags.append("--ipv6")

    command = [
        "rsync",
        "--recursive",
        "--mkpath",
        "--links",
        "--omit-dir-times",
        "--partial",
        "--delay-updates",
        "--times",
        "--temp-dir=/mirror/.debs-incoming/",
        f"--partial-dir={partial_dir}",
        "--human-readable",
        "--progress",
        "--exclude",
        "*_noarch.deb",
        *flags,
        "./debs/",
        upload_url,
    ]

    subprocess.check_call(command)


def have_noarch_files() -> bool:
    output = subprocess.check_output(["find", "./debs", "-name", "*_noarch.deb"])
    return len(output) > 1


def rsync_noarch_file(upload_url: str, branch: str, partial_dir: str, identity_file=None, verbose=False, force_push_noarch_package=False, ipv6=False, ipv4=False, compress=False):
    print("[+] Uploading noarch packages ...")
    flags = []
    ssh_option = ""

    if compress:
        ssh_option = "-C"

    if identity_file:
        flags.extend(["--rsh", f"ssh {ssh_option} -i {identity_file}"])

    if not force_push_noarch_package or branch == "stable":
        flags.append("--ignore-existing")

    if verbose:
        flags.append("--verbose")

    if ipv4:
        flags.append("--ipv4")
    if ipv6:
        flags.append("--ipv6")

    command = [
        "rsync",
        "--recursive",
        "--mkpath",
        "--links",
        "--omit-dir-times",
        "--partial",
        "--delay-updates",
        "--times",
        "--temp-dir=/mirror/.debs-incoming/",
        f"--partial-dir={partial_dir}",
        "--human-readable",
        "--progress",
        "--include",
        "*_noarch.deb",
        *flags,
        "./debs/",
        upload_url,
    ]

    subprocess.check_call(command)


def clean_output_directory():
    print("[+] Cleaning debs ...")
    debs_path = os.path.abspath("./debs")
    subprocess.check_call(["sudo", "rm", "-rv", debs_path])


if __name__ == "__main__":
    main()
