#!/usr/bin/env python
# Copyright 1999-2024 Gentoo Authors
# Distributed under the terms of the GNU General Public License v2

import locale
import os
import sys

if (
    sys.getfilesystemencoding().lower() != "utf-8"
    or locale.getpreferredencoding(False).lower() != "utf-8"
):
    os.environ["PYTHONUTF8"] = "1"
    os.execv(sys.executable, [sys.executable] + sys.argv)


import argparse
import errno
import math
import shlex
import signal
import subprocess

from os import path as osp

if osp.isfile(
    osp.join(osp.dirname(osp.dirname(osp.realpath(__file__))), ".portage_not_installed")
):
    sys.path.insert(
        0, osp.join(osp.dirname(osp.dirname(osp.realpath(__file__))), "lib")
    )
import portage

portage._internal_caller = True
from portage import xpak, gpkg
from portage.const import SUPPORTED_GENTOO_BINPKG_FORMATS
from portage.dbapi.dep_expand import dep_expand
from portage.dep import Atom, use_reduce
from portage.exception import (
    AmbiguousPackageName,
    CompressorOperationFailed,
    InvalidAtom,
    InvalidData,
    InvalidBinaryPackageFormat,
    InvalidDependString,
    PackageSetNotFound,
    PermissionDenied,
)
from portage.util import ensure_dirs, varexpand, _xattr
from portage.util.cpuinfo import makeopts_to_job_count

xattr = _xattr.xattr
from portage._sets import load_default_config, SETPREFIX
from portage.process import find_binary
from portage.util.compression_probe import _compressors
from portage.util._eventloop.global_event_loop import global_event_loop
from portage.gpg import GPG


def quickpkg_atom(options, infos, arg, eout):
    settings = portage.settings
    eroot = portage.settings["EROOT"]
    trees = portage.db[eroot]
    vartree = trees["vartree"]
    vardb = vartree.dbapi
    bintree = trees["bintree"]

    include_config = options.include_config == "y"
    include_unmodified_config = options.include_unmodified_config == "y"
    fix_metadata_keys = ["PF", "CATEGORY"]

    try:
        atom = dep_expand(arg, mydb=vardb, settings=vartree.settings)
    except AmbiguousPackageName as e:
        # Multiple matches thrown from cpv_expand
        eout.eerror(f"Please use a more specific atom: {' '.join(e.args[0])}")
        del e
        infos["missing"].append(arg)
        return 1
    except (InvalidAtom, InvalidData):
        eout.eerror(f"Invalid atom: {arg}")
        infos["missing"].append(arg)
        return 1
    # Check if dep_expand added an operator that wasn't in the original arg
    if atom.operator == "=" and not str(arg).startswith("="):
        # dep_expand() allows missing '=' but it's really invalid
        eout.eerror(f"Invalid atom: {arg}")
        infos["missing"].append(arg)
        return 1

    matches = vardb.match(atom)
    pkgs_for_arg = 0
    retval = 0
    for cpv in matches:
        excluded_config_files = []
        dblnk = vardb._dblink(cpv)
        have_lock = False

        if "__PORTAGE_INHERIT_VARDB_LOCK" not in settings:
            try:
                dblnk.lockdb()
                have_lock = True
            except PermissionDenied:
                pass

        try:
            if not dblnk.exists():
                # unmerged by a concurrent process
                continue
            iuse, use, restrict = vardb.aux_get(cpv, ["IUSE", "USE", "RESTRICT"])
            iuse = [x.lstrip("+-") for x in iuse.split()]
            use = use.split()
            try:
                restrict = use_reduce(restrict, uselist=use, flat=True)
            except InvalidDependString as e:
                eout.eerror(
                    "Invalid RESTRICT metadata " + f"for '{cpv}': {str(e)}; skipping"
                )
                del e
                continue
            if "bindist" in iuse and "bindist" not in use:
                eout.ewarn(f"{cpv}: package was emerged with USE=-bindist!")
                eout.ewarn(f"{cpv}: it might not be legal to redistribute this.")
            elif "bindist" in restrict:
                eout.ewarn(f"{cpv}: package has RESTRICT=bindist!")
                eout.ewarn(f"{cpv}: it might not be legal to redistribute this.")
            eout.ebegin(f"Building package for {cpv}")
            pkgs_for_arg += 1
            existing_metadata = dict(
                zip(fix_metadata_keys, vardb.aux_get(cpv, fix_metadata_keys))
            )
            category, pf = portage.catsplit(cpv)
            required_metadata = {}
            required_metadata["CATEGORY"] = category
            required_metadata["PF"] = pf
            update_metadata = {}
            for k, v in required_metadata.items():
                if v != existing_metadata[k]:
                    update_metadata[k] = v
            if update_metadata:
                vardb.aux_update(cpv, update_metadata)

            binpkg_format = settings.get(
                "BINPKG_FORMAT", SUPPORTED_GENTOO_BINPKG_FORMATS[0]
            )
            if binpkg_format == "xpak":
                xpdata = xpak.xpak(dblnk.dbdir)
                binpkg_tmpfile = os.path.join(
                    bintree.pkgdir, cpv + ".tbz2." + str(portage.getpid())
                )
                ensure_dirs(os.path.dirname(binpkg_tmpfile))
                binpkg_compression = settings.get("BINPKG_COMPRESS", "bzip2")
                try:
                    compression = _compressors[binpkg_compression]
                except KeyError as e:
                    if binpkg_compression:
                        eout.eerror(
                            f"Invalid or unsupported compression method: {e.args[0]}"
                        )
                        return 1
                    # Empty BINPKG_COMPRESS disables compression.
                    binpkg_compression = "none"
                    compression = {
                        "compress": "cat",
                        "package": "sys-apps/coreutils",
                    }

                if (
                    settings.get(
                        f"BINPKG_COMPRESS_FLAGS_{binpkg_compression.upper()}", None
                    )
                    is not None
                ):
                    compression["compress"] = compression["compress"].replace(
                        "${BINPKG_COMPRESS_FLAGS}",
                        f"${{BINPKG_COMPRESS_FLAGS_{binpkg_compression.upper()}}}",
                    )

                try:
                    compression_binary = shlex.split(
                        varexpand(compression["compress"], mydict=settings)
                    )[0]
                except IndexError as e:
                    eout.eerror(
                        f"Invalid or unsupported compression method: {e.args[0]}"
                    )
                    return 1
                if find_binary(compression_binary) is None:
                    missing_package = compression["package"]
                    eout.eerror(
                        f"File compression unsupported {binpkg_compression} (missing package: {missing_package})"
                    )
                    return 1

                cmd = compression["compress"].replace(
                    "{JOBS}",
                    str(makeopts_to_job_count(settings.get("MAKEOPTS", "1"))),
                )
                cmd = shlex.split(varexpand(cmd, mydict=settings))
                # Filter empty elements that make Popen fail
                cmd = [x for x in cmd if x != ""]
                with open(binpkg_tmpfile, "wb") as fobj:
                    proc = subprocess.Popen(cmd, stdin=subprocess.PIPE, stdout=fobj)
                    excluded_config_files = dblnk.quickpkg(
                        proc.stdin,
                        include_config=include_config,
                        include_unmodified_config=include_unmodified_config,
                    )
                    proc.stdin.close()
                    if proc.wait() != os.EX_OK:
                        eout.eend(1)
                        eout.eerror(f"Compressor failed for package {cpv}")
                        retval |= 1
                        try:
                            os.unlink(binpkg_tmpfile)
                        except OSError as e:
                            if e.errno not in (errno.ENOENT, errno.ESTALE):
                                raise
                        continue
                xpak.tbz2(binpkg_tmpfile).recompose_mem(xpdata)
            elif binpkg_format == "gpkg":
                metadata = gpkg.gpkg(settings)._generate_metadata_from_dir(dblnk.dbdir)
                binpkg_tmpfile = os.path.join(
                    bintree.pkgdir, cpv + ".gpkg.tar." + str(os.getpid())
                )
                ensure_dirs(os.path.dirname(binpkg_tmpfile))
                excluded_config_files = dblnk.quickpkg(
                    binpkg_tmpfile,
                    metadata,
                    include_config=include_config,
                    include_unmodified_config=include_unmodified_config,
                )
            else:
                raise InvalidBinaryPackageFormat(binpkg_format)
        except CompressorOperationFailed as e:
            eout.eerror(f"Compressor operation failed")
            os.remove(binpkg_tmpfile)
            binpkg_tmpfile = None
        finally:
            if have_lock:
                dblnk.unlockdb()
        if binpkg_tmpfile is not None:
            pkg_info = bintree.inject(cpv, current_pkg_path=binpkg_tmpfile)
            # The pkg_info value ensures that the following getname call
            # returns the correct path when FEATURES=binpkg-multi-instance
            # is enabled, but fallback to cpv in case the inject call
            # returned None due to some kind of failure.
            binpkg_path = bintree.getname(pkg_info or cpv)
            try:
                s = os.stat(binpkg_path)
            except OSError:
                s = None

        if binpkg_tmpfile is None or s is None or pkg_info is None:
            # Sanity check, shouldn't happen normally.
            eout.eend(1)
            eout.eerror(f"Failed to create package: '{cpv}'")
            retval |= 1
        else:
            eout.eend(0)
            infos["successes"].append((cpv, s.st_size))
            infos["config_files_excluded"] += len(excluded_config_files)
            for filename in excluded_config_files:
                eout.ewarn(f"Excluded config: '{filename}'")
    if not pkgs_for_arg:
        eout.eerror("Could not find anything " + f"to match '{arg}'; skipping")
        infos["missing"].append(arg)
        retval |= 1
    return retval


def quickpkg_set(options, infos, arg, eout):
    eroot = portage.settings["EROOT"]
    trees = portage.db[eroot]
    vartree = trees["vartree"]

    settings = vartree.settings
    settings._init_dirs()
    setconfig = load_default_config(settings, trees)
    sets = setconfig.getSets()

    set_name = arg[1:]
    if not set_name in sets:
        eout.eerror(f"Package set not found: '{arg}'; skipping")
        infos["missing"].append(arg)
        return 1

    try:
        atoms = setconfig.getSetAtoms(set_name)
    except PackageSetNotFound as e:
        eout.eerror(
            f"Failed to process package set '{set_name}' because "
            + f"it contains the non-existent package set '{e}'; skipping"
        )
        infos["missing"].append(arg)
        return 1
    retval = os.EX_OK
    for atom in atoms:
        retval |= quickpkg_atom(options, infos, atom, eout)
    return retval


def quickpkg_extended_atom(options, infos, atom, eout):
    eroot = portage.settings["EROOT"]
    trees = portage.db[eroot]
    vartree = trees["vartree"]
    vardb = vartree.dbapi

    require_metadata = atom.slot or atom.repo
    atoms = []
    for cpv in vardb.cpv_all():
        cpv_atom = Atom(f"={cpv}")

        if atom == "*/*":
            atoms.append(cpv_atom)
            continue

        if not portage.match_from_list(atom, [cpv]):
            continue

        if require_metadata:
            try:
                cpv = vardb._pkg_str(cpv, atom.repo)
            except (KeyError, InvalidData):
                continue
            if not portage.match_from_list(atom, [cpv]):
                continue

        atoms.append(cpv_atom)

    for atom in atoms:
        quickpkg_atom(options, infos, atom, eout)


def quickpkg_main(options, args, eout):
    eroot = portage.settings["EROOT"]
    trees = portage.db[eroot]
    bintree = trees["bintree"]

    try:
        ensure_dirs(bintree.pkgdir)
    except portage.exception.PortageException:
        pass
    if not os.access(bintree.pkgdir, os.W_OK):
        eout.eerror(f"No write access to '{bintree.pkgdir}'")
        return errno.EACCES

    if "xattr" in portage.settings.features and not _xattr.XATTRS_WORKS:
        eout.eerror(
            "No xattr support library was found, " "so xattrs will not be preserved!"
        )
        portage.settings.unlock()
        portage.settings.features.remove("xattr")
        portage.settings.lock()

    infos = {}
    infos["successes"] = []
    infos["missing"] = []
    infos["config_files_excluded"] = 0
    for arg in args:
        if arg[0] == SETPREFIX:
            quickpkg_set(options, infos, arg, eout)
            continue
        try:
            atom = Atom(arg, allow_wildcard=True, allow_repo=True)
        except (InvalidAtom, InvalidData):
            # maybe it's valid but missing category (requires dep_expand)
            quickpkg_atom(options, infos, arg, eout)
        else:
            if atom.extended_syntax:
                quickpkg_extended_atom(options, infos, atom, eout)
            else:
                quickpkg_atom(options, infos, atom, eout)

    if not infos["successes"]:
        eout.eerror("No packages found")
        return 1
    print()
    eout.einfo(f"Packages now in '{bintree.pkgdir}':")
    units = {10: "K", 20: "M", 30: "G", 40: "T", 50: "P", 60: "E", 70: "Z", 80: "Y"}
    for cpv, size in infos["successes"]:
        if not size:
            # avoid OverflowError in math.log()
            size_str = "0"
        else:
            power_of_2 = math.log(size, 2)
            power_of_2 = 10 * (power_of_2 // 10)
            unit = units.get(power_of_2)
            if unit:
                size = float(size) / (2**power_of_2)
                size_str = f"{size:.1f}"
                if len(size_str) > 4:
                    # emulate `du -h`, don't show too many sig figs
                    size_str = str(int(size))
                size_str += unit
            else:
                size_str = str(size)
        eout.einfo(f"{cpv}: {size_str}")
    if infos["config_files_excluded"]:
        print()
        eout.ewarn(f"Excluded config files: {infos['config_files_excluded']}")
        eout.ewarn("See --help if you would like to include config files.")
    if infos["missing"]:
        print()
        eout.ewarn("The following packages could not be found:")
        eout.ewarn(" ".join(infos["missing"]))
        return 2
    return os.EX_OK


if __name__ == "__main__":
    usage = "quickpkg [options] <list of package atoms or package sets>"
    parser = argparse.ArgumentParser(usage=usage)
    parser.add_argument(
        "--umask",
        default="0077",
        help="umask used during package creation (default is 0077)",
    )
    parser.add_argument(
        "--ignore-default-opts",
        action="store_true",
        help="do not use the QUICKPKG_DEFAULT_OPTS environment variable",
    )
    parser.add_argument(
        "--include-config",
        choices=["y", "n"],
        default="n",
        metavar="<y|n>",
        help="include all files protected by CONFIG_PROTECT (as a security precaution, default is 'n')",
    )
    parser.add_argument(
        "--include-unmodified-config",
        choices=["y", "n"],
        default="n",
        metavar="<y|n>",
        help="include files protected by CONFIG_PROTECT that have not been modified since installation (as a security precaution, default is 'n')",
    )
    options, args = parser.parse_known_args(sys.argv[1:])
    if not options.ignore_default_opts:
        default_opts = shlex.split(portage.settings.get("QUICKPKG_DEFAULT_OPTS", ""))
        options, args = parser.parse_known_args(default_opts + sys.argv[1:])
    if not args:
        parser.error("no packages atoms given")
    try:
        umask = int(options.umask, 8)
    except ValueError:
        parser.error(f"invalid umask: {options.umask}")
    # We need to ensure a sane umask for the packages that will be created.
    old_umask = os.umask(umask)
    eout = portage.output.EOutput()

    def sigwinch_handler(signum, frame):
        lines, eout.term_columns = portage.output.get_term_size()

    gpg = None
    if portage.settings.get("BINPKG_GPG_SIGNING_KEY", None):
        gpg = GPG(portage.settings)
        gpg.unlock()

    signal.signal(signal.SIGWINCH, sigwinch_handler)
    try:
        retval = quickpkg_main(options, args, eout)
    finally:
        os.umask(old_umask)
        signal.signal(signal.SIGWINCH, signal.SIG_DFL)
        if gpg is not None:
            gpg.stop()
        global_event_loop().close()
    sys.exit(retval)
