#!/usr/bin/env python3
"""USM system-package-manager helper for APT-based systems.

Implements the USM SPM contract (see usm/README.md, "System package manager
integration") for Debian-style systems:

    usm-spm-apt query <usm-ref>...   -> contract JSON on STDOUT
    usm-spm-apt install <name>...    -> contract JSONL events on STDOUT

File queries are answered from apt-file's Contents index ONLY. This helper
therefore requires apt-file to be installed and its index updated
(`apt-file update`; ~100 MB of Contents per enabled architecture -- the
container bootstrap used by `usm deploy --spm apt` performs both). There
are no package-name heuristics: any ref whose translated path has no owner
in the index lands in not-found. The two sanctioned exceptions are curated
tables: SONAME_PACKAGES (sonames whose Debian package lags the upstream
soname the manifests reference) and PACKAGE_REMAPS (owners remapped to the
package that makes the toolchain usable, e.g. valac-bin to the valac
meta-package carrying the split-out vapis). One
deviation from the DNF helper, which
unions repository filelists with the local rpmdb: files that exist only on
the local system (packages installed from .debs outside the configured
repositories) are invisible to apt-file, so such refs land in not-found.
The DPKG database is consulted for package NAMES only, to compute
installed-dependency-count.

The query translation mirrors spm/dnf with Debian multiarch paths,
preferring x86_64 and following the machine architecture otherwise:

    bin/sbin/libexec -> /usr/bin, /usr/sbin, /usr/bin (Debian ships helper
                        executables under /usr/bin; no /usr/libexec tree)
    lib               -> /usr/lib/<triple>/<soname> and /lib/<triple>/<soname>
                         (usr-merge: both path forms appear across suites)
    gio               -> /usr/lib/<triple>/gio/modules/<module>
    pc                -> /usr/lib/<triple>/pkgconfig/<n> and /usr/share/pkgconfig/<n>
    inc               -> /usr/include/<P> prefix
    typelib           -> /usr/lib/<triple>/girepository-1.0/<n>
    vapi              -> /usr/share/vala/vapi/<n> and /usr/share/vala-*/vapi/<n>
    gir/res/cfg/man/info/locale/rootpath/tag -> as in the DNF table

apt-file is invoked as `apt-file search --fixed-string` (exact paths) and
`apt-file search --substring-match` (anchored prefixes: inc: and the
versioned vapi directory), with all patterns of one mode batched via
--from-file (falling back to one invocation per pattern on apt-file builds
without --from-file). apt-file 3.3 reads --fixed-string as a full-path
exact match and exits 1 without diagnostics when a pattern has no owner,
which this helper treats as not-found rather than failure; every result is
additionally re-validated against an exact or anchored-prefix path
predicate before being credited to a ref, so no over-matches leak into the
contract output.

dependency-count models "a solo install on this machine" exactly like the
DNF helper's repository-only sack: python3-apt resolves the package with
DPkg status pointed at /dev/null (nothing installed), so every dependency
resolves as an install; the overlap with the real DPKG status is reported
as installed-dependency-count. If the python3-apt cache cannot be built or
a package cannot be marked, counts degrade to the DNF helper's fallback
(1, installed ? 1 : 0) with a note on STDERR.

Install runs one `apt-get install -y` transaction with
DEBIAN_FRONTEND=noninteractive and never prompts. An `apt-get --simulate`
pass first provides the transaction size and order for the begin event;
real-run "Get:", "Unpacking" and "Setting up" lines are then mapped to
(coalesced) package progress events and package-complete events. There is
no --test mode: apt cannot run the dpkg transaction phases that the DNF
--test mode exercises. Exit codes match the contract table, with 5 meaning
a dpkg transaction failure (apt's analogue of the rpm code).
"""

import argparse
import json
import os
import platform
import re
import shutil
import subprocess
import sys
import tempfile

import apt.progress.base
import apt_pkg

MULTIARCH_TRIPLES = {
    "x86_64": "x86_64-linux-gnu",
    "aarch64": "aarch64-linux-gnu",
    "armv7l": "arm-linux-gnueabihf",
    "i386": "i386-linux-gnu",
    "i686": "i386-linux-gnu",
    "ppc64le": "powerpc64le-linux-gnu",
    "s390x": "s390x-linux-gnu",
}
TRIPLE = MULTIARCH_TRIPLES.get(platform.machine(), "x86_64-linux-gnu")

# Sonames whose Debian package lags the soname the Web-Stack manifests
# reference: the Contents index has no owner for the exact path, but the
# library is present under an older soname (libsodium23 ships
# libsodium.so.23 while the manifests reference the upstream .26). Curated
# like the emerge helper's atom exception table; entries are validated
# against the apt cache so a stale mapping degrades to not-found instead
# of resolving a ghost package.
SONAME_PACKAGES = {
    "libsodium.so.26": "libsodium23",
}

# Owners the Contents index reports for a file, remapped to the package
# that actually makes the toolchain usable: Debian splits valac's vapis
# into valac-<version>-vapi, which valac-bin does not depend on, and keeps
# g-ir-compiler's real binary in gobject-introspection-bin-linux, which
# gobject-introspection-bin (owner of the dangling /usr/bin symlink) does
# not depend on either; the valac and gobject-introspection meta-packages
# pull each split together. Answering with the meta-package keeps one
# install transaction sufficient.
PACKAGE_REMAPS = {
    "valac-bin": "valac",
    "gobject-introspection-bin": "gobject-introspection",
}

EXIT_OK = 0
EXIT_FAILURE = 1
EXIT_RESOLVE = 3
EXIT_DOWNLOAD = 4
EXIT_TRANSACTION = 5

GET_LINE = re.compile(r"^Get:\d+\s")
UNPACK_LINE = re.compile(r"^Unpacking\s+(\S+)")
SETUP_LINE = re.compile(r"^Setting up\s+(\S+)")
INST_LINE = re.compile(r"^Inst\s+(\S+)")
ERROR_LINE = re.compile(r"^E:\s+")


class HelperError(Exception):
    """Fatal helper failure; message is emitted as a contract error event."""

    def __init__(self, message, exit_code=EXIT_FAILURE):
        super(HelperError, self).__init__(message)
        self.exit_code = exit_code


def emit_event(event):
    """Write one contract JSONL event to STDOUT and flush it."""
    sys.stdout.write(json.dumps(event) + "\n")
    sys.stdout.flush()


def fail(message, exit_code=EXIT_FAILURE):
    """Emit the terminal error event and exit non-zero."""
    emit_event({"type": "error", "message": message})
    sys.exit(exit_code)


def fail_plain(message, exit_code=EXIT_FAILURE):
    """Report a query failure on STDERR only, keeping STDOUT contract-clean."""
    sys.stderr.write("usm-spm-apt: %s\n" % message)
    sys.exit(exit_code)


def plain_name(token):
    """Strip a multiarch :arch qualifier from an apt/dpkg package token."""
    return token.split(":", 1)[0]


def exact(path):
    """Descriptor for one exact indexed path (apt-file --fixed-string)."""
    return ("exact", path, lambda candidate, path=path: candidate == path)


def prefixed(path):
    """Descriptor matching indexed paths under path (--substring-match)."""
    return ("substring", path,
            lambda candidate, path=path: candidate.startswith(path))


def vapi_versioned(resource):
    """Descriptor for /usr/share/vala-*/vapi/<resource>."""
    return (
        "substring",
        "/usr/share/vala-",
        lambda candidate, resource=resource: (
            candidate.startswith("/usr/share/vala-")
            and candidate.endswith("/vapi/" + resource)
        ),
    )


def translate_ref(ref):
    """Map a USM resource ref to apt-file (mode, pattern, predicate) triples.

    Mode is "exact" for --fixed-string full-path matches and "substring"
    for --substring-match anchored prefixes. Returns None for refs whose
    type has no file translation.
    """
    prefix, sep, resource = ref.partition(":")
    if not sep or not resource or any(ord(c) < 32 for c in ref):
        return None
    file_roots = {
        "bin": "/usr/bin",
        "sbin": "/usr/sbin",
        "libexec": "/usr/bin",
        "gir": "/usr/share/gir-1.0",
        "res": "/usr/share",
        "cfg": "/etc",
        "man": "/usr/share/man",
        "info": "/usr/share/info",
        "locale": "/usr/share/locale",
    }
    if prefix in file_roots:
        return [exact(file_roots[prefix] + "/" + resource)]
    if prefix == "typelib":
        return [exact("/usr/lib/%s/girepository-1.0/%s" % (TRIPLE, resource))]
    if prefix == "vapi":
        return [
            exact("/usr/share/vala/vapi/" + resource),
            vapi_versioned(resource),
        ]
    if prefix == "lib":
        return [
            exact("/usr/lib/%s/%s" % (TRIPLE, resource)),
            exact("/lib/%s/%s" % (TRIPLE, resource)),
        ]
    if prefix == "gio":
        return [exact("/usr/lib/%s/gio/modules/%s" % (TRIPLE, resource))]
    if prefix == "pc":
        return [
            exact("/usr/lib/%s/pkgconfig/%s" % (TRIPLE, resource)),
            exact("/usr/share/pkgconfig/" + resource),
        ]
    if prefix == "inc":
        return [prefixed("/usr/include/" + resource)]
    if prefix == "rootpath":
        return [exact("/" + resource.lstrip("/"))]
    if prefix == "tag":
        return [exact("/usr/share/usm-tags/" + resource.replace(".", "/") + ".tag")]
    return None


def first_error_line(text):
    for line in text.splitlines():
        if ERROR_LINE.match(line):
            return line.strip()
    return None


def parse_owner_lines(text):
    """Extract (package, path) pairs from apt-file's "package: path" lines.

    apt-file's progress chatter and anything that does not look like an
    absolute path after "package: " is ignored.
    """
    owners = []
    for line in text.splitlines():
        package, sep, path = line.strip().partition(": ")
        if not sep or not path.startswith("/") or " " in package:
            continue
        owners.append((plain_name(package), path))
    return owners


def run_apt_file(arguments):
    if shutil.which("apt-file") is None:
        raise HelperError(
            "apt-file is not installed; the SPM bootstrap must install it "
            "and run apt-file update", EXIT_RESOLVE)
    return subprocess.run(
        ["apt-file", "search"] + arguments,
        stdout=subprocess.PIPE, stderr=subprocess.PIPE,
        universal_newlines=True,
    )


def option_unsupported(stderr):
    lowered = stderr.lower()
    return any(marker in lowered for marker in
               ("unknown option", "unrecognized", "usage:", "usage error"))


def apt_file_search(mode, patterns):
    """Search the apt-file Contents index for every pattern in one mode."""
    results = []
    if not patterns:
        return results
    flags = ["--fixed-string"] if mode == "exact" else ["--substring-match"]
    handle = tempfile.NamedTemporaryFile(
        "w", suffix=".usm-spm-apt-patterns", delete=False)
    try:
        handle.write("\n".join(patterns) + "\n")
        handle.close()
        process = run_apt_file(flags + ["--from-file", handle.name])
        if process.returncode == 0:
            return parse_owner_lines(process.stdout)
        if option_unsupported(process.stderr):
            for pattern in patterns:
                process = run_apt_file(flags + [pattern])
                if process.returncode != 0 and first_error_line(
                        process.stderr + process.stdout) is not None:
                    raise HelperError(
                        first_error_line(process.stderr + process.stdout),
                        EXIT_RESOLVE)
                results.extend(parse_owner_lines(process.stdout))
            return results
        error = first_error_line(process.stderr + process.stdout)
        if error is not None:
            raise HelperError(error, EXIT_RESOLVE)
        return results
    finally:
        os.unlink(handle.name)


def dpkg_installed_names():
    """Names of every installed package, from the dpkg status database."""
    try:
        process = subprocess.run(
            ["dpkg-query", "--show", "--showformat", "${Package}\n"],
            stdout=subprocess.PIPE, stderr=subprocess.PIPE,
            universal_newlines=True)
    except OSError:
        return set()
    if process.returncode != 0:
        return set()
    return set(process.stdout.split())


def build_solo_cache():
    """An apt cache whose dpkg status is empty (nothing installed).

    The DNF helper builds its solo-install estimates against a repository-
    only sack; pointing Dir::State::status at /dev/null is apt's analogue,
    so every dependency of a marked package resolves as an install.
    """
    apt_pkg.init()
    apt_pkg.config.set("Dir::State::status", os.devnull)
    return apt_pkg.Cache(apt.progress.base.OpProgress())


def solo_closure(cache, installed_names, name, closure_cache):
    """(dependency-count, installed-dependency-count) for a solo install."""
    if name in closure_cache:
        return closure_cache[name]

    counts = (1, 1 if name in installed_names else 0)
    if cache is None:
        sys.stderr.write(
            "usm-spm-apt: no python3-apt cache available, estimating "
            "dependency counts of %s as installed-only\n" % name)
        closure_cache[name] = counts
        return counts

    try:
        package = cache[name]
    except KeyError:
        package = None
    if package is None:
        sys.stderr.write(
            "usm-spm-apt: %s not found in enabled repositories, "
            "estimating dependency counts as installed-only\n" % name)
        closure_cache[name] = counts
        return counts

    depcache = apt_pkg.DepCache(cache)
    depcache.init()
    if depcache.get_candidate_ver(package) is None:
        sys.stderr.write(
            "usm-spm-apt: %s has no installation candidate, "
            "estimating dependency counts as installed-only\n" % name)
        closure_cache[name] = counts
        return counts

    try:
        depcache.mark_install(package, True, True)
        closure = [p.name for p in cache.packages if depcache.marked_install(p)]
        counts = (
            len(closure),
            sum(1 for member in closure if member in installed_names),
        )
    except SystemError as e:
        sys.stderr.write(
            "usm-spm-apt: could not resolve solo install of %s: %s\n"
            % (name, e))
    closure_cache[name] = counts
    return counts


def soname_exception_package(soname):
    """The cached package for a {@link SONAME_PACKAGES} soname, or None.

    The mapping is only trusted when the apt cache actually offers the
    package, so a stale table entry degrades to not-found.
    """
    name = SONAME_PACKAGES.get(soname)
    if name is None:
        return None
    try:
        cache = apt.Cache()
        if name in cache and cache[name].candidate is not None:
            return name
    except Exception:
        pass
    return None


def cached_package_name(name):
    """The remapped {@link PACKAGE_REMAPS} target for an owner, or None.

    The remap is only trusted when the apt cache actually offers the
    target, so a stale table entry keeps the index-reported owner.
    """
    target = PACKAGE_REMAPS.get(name)
    if target is None:
        return None
    try:
        cache = apt.Cache()
        if target in cache and cache[target].candidate is not None:
            return target
    except Exception:
        pass
    return None


def run_query(args):
    refs = list(dict.fromkeys(args.refs))
    translated = {}
    patterns = {"exact": set(), "substring": set()}
    for ref in refs:
        descriptors = translate_ref(ref)
        if descriptors is not None:
            translated[ref] = descriptors
            for mode, pattern, _ in descriptors:
                patterns[mode].add(pattern)

    index = []
    for mode in ("exact", "substring"):
        index.extend(apt_file_search(mode, sorted(patterns[mode])))

    not_found = []
    candidates = {}
    for ref in refs:
        descriptors = translated.get(ref)
        if descriptors is None:
            sys.stderr.write(
                "usm-spm-apt: resource type of \"%s\" has no "
                "system-package-manager translation\n" % ref)
            not_found.append(ref)
            continue
        matched = set()
        for mode, pattern, predicate in descriptors:
            for name, path in index:
                if predicate(path):
                    matched.add(name)
        remapped = set()
        for name in sorted(matched):
            remapped.add(cached_package_name(name) or name)
        matched = remapped
        if not matched and ref.startswith("lib:"):
            mapped = soname_exception_package(ref[len("lib:"):])
            if mapped is not None:
                matched.add(mapped)
        if not matched:
            not_found.append(ref)
            continue
        for name in matched:
            entry = candidates.get(name)
            if entry is None:
                entry = {"resources": []}
                candidates[name] = entry
            entry["resources"].append(ref)

    installed_names = dpkg_installed_names()
    solo_cache = build_solo_cache() if candidates else None
    closure_cache = {}
    packages = []
    for name in sorted(candidates):
        dependency_count, installed_dependency_count = solo_closure(
            solo_cache, installed_names, name, closure_cache)
        packages.append({
            "name": name,
            "resources": candidates[name]["resources"],
            "dependency-count": dependency_count,
            "installed-dependency-count": installed_dependency_count,
        })

    sys.stdout.write(json.dumps({
        "not-found": not_found,
        "packages": packages,
    }) + "\n")
    sys.stdout.flush()
    return EXIT_OK


def apt_environment():
    environment = dict(os.environ)
    environment["DEBIAN_FRONTEND"] = "noninteractive"
    environment["LC_ALL"] = "C"
    return environment


def classify_apt_failure(output):
    lowered = output.lower()
    if "failed to fetch" in lowered or "unable to fetch" in lowered:
        return EXIT_DOWNLOAD
    if ("unable to locate" in lowered or "unmet dependencies" in lowered
            or "no installation candidate" in lowered):
        return EXIT_RESOLVE
    return EXIT_TRANSACTION


def collect_error_messages(output, limit=5):
    messages = [line.strip() for line in output.splitlines()
                if ERROR_LINE.match(line)]
    return "; ".join(messages[-limit:])


def simulate_install(names, environment):
    """Ordered package list of the would-be transaction, from -s output."""
    process = subprocess.run(
        ["apt-get", "--simulate", "install"] + names,
        stdout=subprocess.PIPE, stderr=subprocess.PIPE,
        universal_newlines=True, env=environment)
    if process.returncode != 0:
        combined = process.stderr + process.stdout
        raise HelperError(
            collect_error_messages(combined)
            or "apt-get simulation failed for %s" % " ".join(names),
            classify_apt_failure(combined))
    order = []
    for line in process.stdout.splitlines():
        match = INST_LINE.match(line)
        if match:
            name = plain_name(match.group(1))
            if name not in order:
                order.append(name)
    return order


def cmd_plan(args):
    environment = apt_environment()
    try:
        order = simulate_install(args.names, environment)
    except HelperError as e:
        fail_plain(str(e), e.exit_code)
    sys.stdout.write(json.dumps({"packages": order}) + "\n")
    sys.stdout.flush()
    return EXIT_OK


def cmd_query(args):
    try:
        return run_query(args)
    except SystemExit:
        raise
    except HelperError as e:
        fail_plain(str(e), e.exit_code)
    except Exception as e:
        fail_plain("query failed: %s" % e, EXIT_RESOLVE)


class InstallProgress:
    """Maps apt-get output lines to contract package/complete events."""

    def __init__(self, order):
        self.position = {name: index for index, name in enumerate(order, 1)}
        self.total = len(order)
        self.completed = 0
        self.last = None

    def emit_package(self, name, progress):
        if name not in self.position:
            return
        event = ("package", name, self.position[name], self.total,
                 round(progress, 4))
        if event == self.last:
            return
        self.last = event
        emit_event({
            "type": "package",
            "name": name,
            "current": self.position[name],
            "total": self.total,
            "progress": round(progress, 4),
        })

    def feed(self, line):
        match = GET_LINE.match(line)
        if match:
            for token in line.split()[1:]:
                candidate = plain_name(token)
                if candidate in self.position:
                    self.emit_package(candidate, 0.0)
                    break
            return
        match = UNPACK_LINE.match(line)
        if match:
            self.emit_package(plain_name(match.group(1)), 0.5)
            return
        match = SETUP_LINE.match(line)
        if match:
            name = plain_name(match.group(1))
            if name in self.position:
                self.emit_package(name, 1.0)
                emit_event({"type": "package-complete", "name": name})
                self.completed += 1


def cmd_install(args):
    environment = apt_environment()
    try:
        order = simulate_install(args.names, environment)
    except HelperError as e:
        fail(str(e), e.exit_code)

    emit_event({"type": "begin", "total": len(order)})
    progress = InstallProgress(order)

    process = subprocess.Popen(
        ["apt-get", "install", "-y"] + args.names,
        stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
        universal_newlines=True, env=environment)
    output = []
    for line in process.stdout:
        output.append(line)
        progress.feed(line.rstrip("\r\n"))
    process.stdout.close()
    returncode = process.wait()

    transcript = "".join(output)
    if returncode != 0:
        fail(collect_error_messages(transcript)
             or "apt-get install failed with exit code %d" % returncode,
             classify_apt_failure(transcript))

    emit_event({
        "type": "complete",
        "status": "ok",
        "installed": progress.completed,
    })
    return EXIT_OK


def main():
    parser = argparse.ArgumentParser(
        prog="usm-spm-apt",
        description="USM system-package-manager helper for APT")
    subparsers = parser.add_subparsers(dest="command", required=True)

    query_parser = subparsers.add_parser(
        "query", help="resolve USM resource refs to system packages")
    query_parser.add_argument(
        "refs", nargs="+", metavar="USM-REF",
        help="resource ref, e.g. bin:valac or lib:libglib-2.0.so.0")
    query_parser.set_defaults(handler=cmd_query)

    plan_parser = subparsers.add_parser(
        "plan", help="resolve the material install set for names "
                     "(chosen packages plus all dependencies)")
    plan_parser.add_argument(
        "names", nargs="+", metavar="NAME",
        help="native system package name")
    plan_parser.set_defaults(handler=cmd_plan)

    install_parser = subparsers.add_parser(
        "install", help="install system packages, streaming progress events")
    install_parser.add_argument(
        "names", nargs="+", metavar="NAME",
        help="native system package name")
    install_parser.set_defaults(handler=cmd_install)

    args = parser.parse_args()
    try:
        return args.handler(args)
    except HelperError as e:
        fail(str(e), e.exit_code)
    except SystemExit:
        raise
    except Exception as e:
        fail("unexpected failure: %s" % e, EXIT_FAILURE)


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