#!/usr/bin/env python3

from collections import namedtuple
import sys
import os
import os.path
import re
import subprocess
import tempfile
import urllib.request
import argparse
import hashlib

DistroInfo = namedtuple('DistroInfo',
                        ['sources_list', 'url_keys', 'public_keys'])

_DISTROS = {
    "ubuntu":
    DistroInfo(
        """deb http://us.archive.ubuntu.com/ubuntu/ {release} main restricted
deb-src http://us.archive.ubuntu.com/ubuntu/ {release} main restricted

deb http://security.ubuntu.com/ubuntu {release}-security main restricted
deb-src http://security.ubuntu.com/ubuntu {release}-security main restricted

deb http://us.archive.ubuntu.com/ubuntu/ {release} universe
deb-src http://us.archive.ubuntu.com/ubuntu/ {release} universe
deb http://us.archive.ubuntu.com/ubuntu/ {release}-updates universe
deb-src http://us.archive.ubuntu.com/ubuntu/ {release}-updates universe

deb http://us.archive.ubuntu.com/ubuntu/ {release}-updates main restricted
deb-src http://us.archive.ubuntu.com/ubuntu/ {release}-updates main restricted

deb http://us.archive.ubuntu.com/ubuntu {release}-backports main restricted
deb-src http://us.archive.ubuntu.com/ubuntu {release}-backports main restricted""",
        [], ['3B4FE6ACC0B21F32', '871920D1991BC93C']),
    "debian":
    DistroInfo(
        """deb http://deb.debian.org/debian/ {release} main contrib non-free
deb-src http://deb.debian.org/debian/ {release} main contrib non-free

deb https://security.debian.org/debian-security {release}-security main contrib non-free
deb-src https://security.debian.org/debian-security {release}-security main contrib non-free

deb http://deb.debian.org/debian/ {release}-updates main contrib non-free
deb-src http://deb.debian.org/debian/ {release}-updates main contrib non-free

deb http://deb.debian.org/debian {release}-backports main contrib non-free
deb-src http://deb.debian.org/debian {release}-backports main contrib non-free""",
        [
            "https://ftp-master.debian.org/keys/archive-key-11.asc",
            "https://ftp-master.debian.org/keys/archive-key-11-security.asc"
        ], []),
}


def initialize_apt(apt_dir, apt_args, args):
    os.mkdir(os.path.join(apt_dir, 'etc'))
    os.mkdir(os.path.join(apt_dir, 'etc', 'apt'))
    os.mkdir(os.path.join(apt_dir, 'etc', 'apt', 'trusted.gpg.d'))
    os.mkdir(os.path.join(apt_dir, 'etc', 'apt', 'preferences.d'))
    os.mkdir(os.path.join(apt_dir, 'var'))
    os.mkdir(os.path.join(apt_dir, 'var', 'lib'))
    os.mkdir(os.path.join(apt_dir, 'var', 'lib', 'dpkg'))
    with open(os.path.join(apt_dir, 'var', 'lib', 'dpkg', 'status'), 'w'):
        pass
    distro = _DISTROS[args.distro]
    with open(os.path.join(apt_dir, 'etc', 'apt', 'sources.list'), 'w') as f:
        f.write(distro.sources_list.format(release=args.release))
    for key in args.apt_key + distro.url_keys:
        basename = os.path.basename(key)
        urllib.request.urlretrieve(
            key, os.path.join(apt_dir, 'etc', 'apt', 'trusted.gpg.d',
                              basename))
    for key in distro.public_keys:
        subprocess.check_call([
            "apt-key", "--keyring",
            os.path.join(apt_dir, 'etc', 'apt', 'trusted.gpg'), "adv",
            "--keyserver", "keyserver.ubuntu.com", "--recv-key", key
        ])

    subprocess.check_call(["apt-get"] + apt_args + ["update"])


def get_deps(apt_args, package):
    env = dict(os.environ)
    if 'LD_LIBRARY_PATH' in env:
        del env['LD_LIBRARY_PATH']
    out = subprocess.check_output(["apt-rdepends"] + apt_args + [package],
                                  env=env)
    deps = out.splitlines()
    return set([dep for dep in deps if not dep.startswith(b" ")])


def get_all_deps(apt_args, packages):
    deps = set()
    for package in packages or ():
        deps.update(get_deps(apt_args, package))
    return deps


def map_virtual_packages(packages):
    '''Maps known virtual packages to the preferred concrete packages which
  provide them.'''
    for package in packages:
        if package == b'python-numpy-abi9':
            yield b'python-numpy'
            continue
        if package == b'python3-numpy-abi9':
            yield b'python3-numpy'
            continue
        if package == b'libjack-0.125':
            yield b'libjack-jackd2-0'
            continue
        if package == b'fonts-freefont':
            yield b'fonts-freefont-ttf'
            continue
        if package == b'gsettings-backend':
            yield b'dconf-gsettings-backend'
            continue
        if package == b'gdal-abi-2-4-0':
            yield b'libgdal20'
            continue
        if package == b'libglu1':
            yield b'libglu1-mesa'
            continue
        if package == b'liblapack.so.3':
            yield b'liblapack3'
            continue
        if package == b'libopencl1':
            yield b'ocl-icd-libopencl1'
            continue
        if package == b'libgcc1':
            yield b'libgcc-s1'
            continue
        if package == b'libopencl-1.2-1':
            yield b'ocl-icd-libopencl1'
            continue
        if package == b'libblas.so.3':
            yield b'libblas3'
            continue
        if package == b'debconf-2.0':
            yield b'debconf'
            continue
        yield package


def download_deps(apt_args, packages, excludes, force_includes,
                  force_excludes):
    deps = get_all_deps(apt_args, packages)
    exclude_deps = get_all_deps(apt_args, excludes)
    exclude_deps -= set({b'libssl1.1'})
    deps -= exclude_deps
    force_include_deps = get_all_deps(apt_args, force_includes)
    deps |= force_include_deps
    force_exclude_deps = get_all_deps(apt_args, force_excludes)
    deps -= force_exclude_deps
    env = dict(os.environ)
    if 'LD_LIBRARY_PATH' in env:
        del env['LD_LIBRARY_PATH']
    subprocess.check_call([b"apt-get"] + [a.encode('utf-8')
                                          for a in apt_args] + [b"download"] +
                          list(map_virtual_packages(deps)),
                          env=env)


def fixup_files():
    # Gotta remove those pesky epoch numbers in the file names. Bazel doesn't
    # like them.
    regex = re.compile(".%3a")
    contents = os.listdir(os.getcwd())
    for deb in contents:
        new_name = regex.sub("", deb)
        if new_name != deb:
            os.rename(deb, new_name)


def sha256_checksum(filename, block_size=65536):
    sha256 = hashlib.sha256()
    with open(filename, 'rb') as f:
        for block in iter(lambda: f.read(block_size), b''):
            sha256.update(block)
    return sha256.hexdigest()


def print_file_list():
    contents = os.listdir(os.getcwd())
    contents.sort()
    print("_files = {")
    for deb in contents:
        print('    "%s": "%s",' % (deb, sha256_checksum(deb)))
    print("}")


_ALWAYS_EXCLUDE = [
    "dbus-session-bus",
    "debconf",
    "debconf-2.0",
    "default-dbus-session-bus",
    "dpkg",
    "install-info",
    "libc-dev",
    "libc6",
    "libc6-dev",
]


def main(argv):
    parser = argparse.ArgumentParser()
    parser.add_argument("--exclude",
                        "-e",
                        type=str,
                        action="append",
                        help="A package to exclude from the list")
    parser.add_argument(
        "--force-include",
        type=str,
        action="append",
        help=
        "Force include this and its dependencies. Even if listed in excludes.")
    parser.add_argument(
        "--force-exclude",
        type=str,
        action="append",
        help=
        "Force exclude this and its dependencies. Even if listed via --force-include."
    )
    parser.add_argument("--arch",
                        type=str,
                        default="amd64",
                        help="Architecture to download files for.")
    parser.add_argument("--distro",
                        type=str,
                        default="debian",
                        choices=_DISTROS.keys(),
                        help="Architecture to download files for.")
    parser.add_argument(
        "--apt-dir",
        type=str,
        help=" ".join([
            "File to generate and store apt files in.",
            "Helpful for saving time when downloading multiple groups of packages.",
            "Some flags will be ignored in favor of the values used to create this folder, so be careful.",
        ]))
    parser.add_argument("--release",
                        type=str,
                        default="bookworm",
                        help="Debian/Ubuntu release to use.")
    parser.add_argument("--apt-key",
                        type=str,
                        action="append",
                        default=[],
                        help="URL of an additional apt archive key to trust.")
    parser.add_argument("package", nargs="+", help="The packages to download.")
    args = parser.parse_args(argv[1:])
    if args.apt_dir:
        apt_dir = args.apt_dir
    else:
        apt_dir = tempfile.mkdtemp()
    apt_args = ["-o", "Dir=" + apt_dir, "-o", "APT::Architecture=" + args.arch]
    if not args.apt_dir:
        print("Creating apt files in %s" % apt_dir)
        initialize_apt(apt_dir, apt_args, args)
    folder = tempfile.mkdtemp()
    os.chdir(folder)
    excludes = args.exclude or []
    # Exclude common packages that don't make sense to include in everything all
    # the time.
    excludes += _ALWAYS_EXCLUDE
    download_deps(apt_args, args.package, excludes, args.force_include,
                  args.force_exclude)
    fixup_files()
    print_file_list()
    print("Your packages are all in %s" % folder)


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