#!/usr/bin/env python3

import argparse
import os
import sys

import portage


def wanted_distfiles():
    vartree = portage.db["/"]["vartree"]
    files = set()
    for name in vartree.getallnodes():
        for package in vartree.dep_match(name):
            try:
                files.update(portage.portdb.getFetchMap(package).keys())
            except Exception as exc:
                print(f"Failed to get file list for {package}: {exc}")
    return files


def main(argv):
    parser = argparse.ArgumentParser(
        prog="go-distclean", description="Remove distfiles no installed package needs."
    )
    parser.add_argument("-p", "--pretend", action="store_true")
    args = parser.parse_args(argv[1:])

    distdir = portage.settings["DISTDIR"]
    print("DISTDIR =", distdir)

    files = wanted_distfiles()
    if not files:
        sys.exit("No package files found. This cannot be right.\n")

    size = 0
    count = 0
    for filename in os.listdir(distdir):
        path = os.path.join(distdir, filename)
        if not os.path.isfile(path) or os.path.islink(path) or filename in files:
            continue
        size += os.stat(path).st_size
        count += 1
        if args.pretend:
            print("Would remove", path)
        else:
            try:
                os.remove(path)
                print("Removed", path)
            except OSError:
                print("Failed to remove", path)

    print("%i files, total size: %i MB" % (count, size // 1048576))
    return 0


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