#!/usr/bin/env python

import json
import os
import subprocess
import sys
import urllib.request
from argparse import ArgumentParser

PYPI_JSON_TEMPLATE = "https://pypi.org/pypi/{package}/{version}/json"


def download_package_files(package, version, dest):
    json_url = PYPI_JSON_TEMPLATE.format(
        package=package,
        version=version,
    )
    with urllib.request.urlopen(json_url) as resp:
        data = json.load(resp)["urls"]

    for entry in data:
        path = os.path.join(dest, entry["filename"])
        print(f"Downloading {entry['url']}")
        with urllib.request.urlopen(entry["url"]) as resp:
            with open(path, "wb") as fh:
                fh.write(resp.read())

        subprocess.run(
            ["sha256sum", "--check"],
            input=f"{entry['digests']['sha256']} {path}".encode("utf-8"),
            check=True,
        )


def download_packages(packages, dest):
    for package in packages:
        try:
            name, version = package.split("==")
            name = name.strip()
            version = version.strip()
        except ValueError:
            print(f"error: {package} must be of the form <name>==<version>")
            sys.exit(1)

        download_package_files(name, version, dest)


def cli(args=sys.argv[1:]):
    parser = ArgumentParser()
    parser.add_argument(
        "packages",
        nargs="+",
        help="One or more packages of the form '<name>==<version>'.",
    )
    parser.add_argument(
        "-d",
        "--dest",
        default=None,
        help="Destination to download package files. Defaults to current working directory.",
    )
    args = parser.parse_args(args)

    dest = args.dest or os.getcwd()
    download_packages(args.packages, dest)


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