#!/usr/bin/env python3
"""Archive a GitHub user's starred repositories by cloning them locally."""

import argparse
import json
import os
import shutil
import subprocess
import sys
import time
from concurrent.futures import ThreadPoolExecutor, as_completed

import requests

SA_VERSION = "1.0.0"
DEFAULT_PER_PAGE = 100
API_URL = "https://api.github.com"
USER_AGENT = "stararchiver/{}".format(SA_VERSION)


def log(msg):
    print(msg, flush=True)


def make_session(token):
    session = requests.Session()
    session.headers.update({
        "User-Agent": USER_AGENT,
        "Accept": "application/vnd.github+json",
    })
    if token:
        session.headers["Authorization"] = "Bearer {}".format(token)
    return session


def iter_starred(username, token, per_page=DEFAULT_PER_PAGE, max_repos=None):
    """Yield starred repositories, following GitHUb API pagination."""
    session = make_session(token)
    url = "{}/users/{}/starred".format(API_URL, username)
    params = {"per_page": per_page, "page": 1}
    seen = 0

    while url:
        resp = session.get(url, params=params)
        params = None
        handle_rate_limit(resp)
        if resp.status_code == 404:
            raise SystemExit(
                "GitHub user '{}' not found (404). "
                "Double-check the username.".format(username)
            )
        if resp.status_code == 401:
            raise SystemExit("GitHub API rejected the token (401). Is it valid?")
        if resp.status_code != 200:
            raise SystemExit(
                "GitHub API error {}: {}".format(resp.status_code, resp.text)
            )

        for repo in resp.json():
            yield repo
            seen += 1
            if max_repos and seen >= max_repos:
                return

        url = next_link(resp.headers.get("Link", ""))


def next_link(link_header):
    for part in link_header.split(","):
        sections = [s.strip() for s in part.split(";")]
        if len(sections) == 2 and sections[1] == 'rel="next"':
            return sections[0].strip("<>")
    return None


def handle_rate_limit(resp):
    remaining = resp.headers.get("X-RateLimit-Remaining")
    retry_after = resp.headers.get("Retry-After")
    is_rate_limited = (
        (resp.status_code in (403, 429) and retry_after is not None)
        or remaining == "0"
    )
    if not is_rate_limited:
        return

    if retry_after is not None:
        wait = int(retry_after)
    else:
        reset = int(resp.headers.get("X-RateLimit-Reset", 0))
        wait = max(reset - int(time.time()), 0)
    wait = min(wait, 3600)
    log(
        "GitHub API rate limit reached; sleeping {}s. "
        "Pass --token or set GITHUB_TOKEN for a higher limit.".format(wait)
    )
    time.sleep(wait)


def clone_repo(dest, url, shallow):
    if os.path.isdir(os.path.join(dest, ".git")):
        return True, None
    if os.path.exists(dest):
        shutil.rmtree(dest)
    cmd = ["git", "clone", "--quiet"]
    if shallow:
        cmd += ["--depth", "1"]
    cmd += [url, dest]
    # Never prompt for credentials: a 401/404 from GitHub (e.g. a wiki that
    # was never created) must fail fast instead of hanging on a login prompt.
    env = dict(os.environ, GIT_TERMINAL_PROMPT="0")
    proc = subprocess.run(cmd, capture_output=True, text=True, env=env)
    if proc.returncode != 0:
        return False, proc.stderr.strip()
    return True, None


def build_clone_url(repo, use_ssh, for_wiki=False):
    url = repo["ssh_url"] if use_ssh else repo["clone_url"]
    if for_wiki:
        url = url[:-4] + ".wiki.git" if url.endswith(".git") else url + ".wiki.git"
    return url


def archive_repo(repo, dest_dir, use_ssh, shallow, with_wiki):
    """Clone a repo and (sequentially) its wiki, so the two never race."""
    results = []
    author = repo["owner"]["login"]
    name = repo["name"]
    repo_dir = os.path.join(dest_dir, "{}-{}".format(author, name))
    label = "{}-{}".format(author, name)

    success, error = clone_repo(repo_dir, build_clone_url(repo, use_ssh), shallow)
    results.append((label, success, error))
    if success and with_wiki and repo.get("has_wiki"):
        wiki_dir = os.path.join(repo_dir, "_wiki")
        wiki_label = "{} (wiki)".format(label)
        wiki_ok, wiki_err = clone_repo(
            wiki_dir, build_clone_url(repo, use_ssh, for_wiki=True), shallow
        )
        results.append((wiki_label, wiki_ok, wiki_err))
    return results


def clone_starred(reports, dest_dir, use_ssh, shallow, with_wiki, jobs):
    os.makedirs(dest_dir, exist_ok=True)
    clones = ok = fails = 0

    with ThreadPoolExecutor(max_workers=max(jobs, 1)) as pool:
        futures = [
            pool.submit(archive_repo, repo, dest_dir, use_ssh, shallow, with_wiki)
            for repo in reports
        ]
        for i, fut in enumerate(as_completed(futures), 1):
            for label, success, error in fut.result():
                clones += 1
                if success:
                    ok += 1
                    log("[{}/{}] ok: {}".format(i, clones, label))
                else:
                    fails += 1
                    log("[{}/{}] FAILED: {}: {}".format(i, clones, label, error))

    log("Cloned {}/{} repo(s) successfully ({} failure(s)).".format(
        ok, clones, fails))
    return dest_dir


def compress_repos(path):
    archive_path = shutil.make_archive(
        os.path.abspath(path), "zip",
        root_dir=os.path.dirname(os.path.abspath(path)),
        base_dir=os.path.basename(os.path.abspath(path)),
    )
    log("Zip archive created: {}".format(archive_path))
    return archive_path


def write_metadata(dest_dir, reports):
    data = []
    for repo in reports:
        data.append({
            "full_name": repo["full_name"],
            "description": repo.get("description"),
            "html_url": repo.get("html_url"),
            "clone_url": repo.get("clone_url"),
            "ssh_url": repo.get("ssh_url"),
            "stars": repo.get("stargazers_count"),
            "language": repo.get("language"),
            "fork": repo.get("fork"),
            "archived": repo.get("archived"),
        })
    path = os.path.join(dest_dir, "starred_repos.json")
    with open(path, "w") as fh:
        json.dump(data, fh, indent=2)
    log("Wrote {} repo(s) metadata to {}".format(len(data), path))
    return path


def get_dest_dir(base_dir, username):
    base = os.path.abspath(os.path.expanduser(base_dir or os.getcwd()))
    return os.path.join(base, "stararchiver_{}_{}".format(username, int(time.time())))


def assert_git():
    if shutil.which("git") is None:
        raise SystemExit(
            "Error: stararchiver requires git to be installed "
            "(e.g. 'brew install git')."
        )


def build_parser():
    parser = argparse.ArgumentParser(
        description="Archive a GitHub user's starred repositories by cloning "
                    "them to disk (and optionally zipping them up).",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument(
        "username", nargs="?", help="GitHub username to archive",
    )
    parser.add_argument(
        "-o", "--output-dir",
        help="Base directory to write archives into (default: current directory)",
    )
    parser.add_argument(
        "-c", "--compress", action="store_true",
        help="Compress the archived repositories into a single zip file",
    )
    parser.add_argument(
        "-d", "--delete-data", action="store_true",
        help="Delete cloned repositories after archiving/compressing",
    )
    parser.add_argument(
        "-t", "--token", default=os.environ.get("GITHUB_TOKEN"),
        help="GitHub API token (or set the GITHUB_TOKEN environment variable)",
    )
    parser.add_argument(
        "--ssh", action="store_true",
        help="Clone using SSH URLs instead of HTTPS (requires SSH keys on GitHub)",
    )
    parser.add_argument(
        "--no-wiki", action="store_true",
        help="Do not clone git wikis",
    )
    parser.add_argument(
        "--shallow", action="store_true",
        help="Shallow-clone (--depth 1) to save disk space and time",
    )
    parser.add_argument(
        "-j", "--jobs", type=int, default=4,
        help="Number of concurrent clones",
    )
    parser.add_argument(
        "--max-repos", type=int, default=None,
        help="Only archive the first N starred repos (useful for testing)",
    )
    parser.add_argument(
        "--metadata", action="store_true",
        help="Write a starred_repos.json manifest of all starred repos",
    )
    parser.add_argument(
        "--dry-run", action="store_true",
        help="List the starred repositories without cloning anything",
    )
    parser.add_argument(
        "-v", "--version", action="store_true", help="Show version and exit",
    )
    return parser


def main(argv=None):
    args = build_parser().parse_args(argv)
    if args.version:
        log("stararchiver {}".format(SA_VERSION))
        return 0
    if args.username is None:
        raise SystemExit(parser.error("the following arguments are required: username"))
    assert_git()

    if not args.token:
        log(
            "No token provided; the unauthenticated GitHub API rate limit "
            "(60 requests/hr) applies. Use --token or set GITHUB_TOKEN for "
            "5000/hr."
        )
    log("Fetching starred repositories for {}...".format(args.username))
    reports = list(iter_starred(args.username, args.token, max_repos=args.max_repos))

    if args.dry_run:
        for repo in reports:
            log("* {}".format(repo["full_name"]))
        log("Found {} starred repository(ies).".format(len(reports)))
        return 0

    log("Found {} starred repository(ies).".format(len(reports)))
    if not reports:
        log("Nothing to archive.")
        return 0

    dest_dir = get_dest_dir(args.output_dir, args.username)
    clone_starred(
        reports,
        dest_dir,
        use_ssh=args.ssh,
        shallow=args.shallow,
        with_wiki=not args.no_wiki,
        jobs=args.jobs,
    )
    if args.metadata:
        write_metadata(dest_dir, reports)
    if args.compress:
        log("Compressing repositories...")
        compress_repos(dest_dir)
    if args.delete_data:
        shutil.rmtree(dest_dir)
        log("Deleted {}.".format(dest_dir))
        log("Done.")
    else:
        log("Done. Archived data is in {}".format(dest_dir))
    return 0


if __name__ == "__main__":
    sys.exit(main())