#!/usr/bin/env python3
"""Work out which Rust crates lightning-matrix must vendor.

Policy is to build against the archive's librust-*-dev packages and vendor
only what the archive cannot provide. This script, run on the target series
(apt must list its archive), finds that set for the current source:

  1. applies debian/patches to a scratch copy and runs `cargo vendor` on it;
  2. unpacks every archive version of each crate it needs (no root needed);
  3. resolves with archive crates first and vendors only what resolution
     requires, dropping every vendored crate that turns out to be unnecessary;
  4. proves the Build-Depends closure covers the archive crates in use.

Writes debian/vendored-crates.txt (input for debian/vendor-crates) and
debian/cargo-build-depends.txt (paste into debian/control). Needs cargo,
network access and apt-get; takes a few minutes.
"""
import collections
import json
import os
import re
import shutil
import subprocess
import sys
import tempfile

TOP = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
WORK = tempfile.mkdtemp(prefix="vendor-set-")
ENV = dict(os.environ, CARGO_HOME=os.path.join(WORK, "cargo-home"))
SRC = os.path.join(WORK, "src")
FULL = os.path.join(WORK, "full")
DEBS = os.path.join(WORK, "debs")
ARCH = os.path.join(WORK, "arch")
MERGED = os.path.join(WORK, "merged")
REGISTRY = "usr/share/cargo/registry"


def run(cmd, **kw):
    return subprocess.run(cmd, capture_output=True, text=True, **kw)


def ser(v):
    return tuple(int(x) for x in re.findall(r"\d+", v)[:3])


def prepare_source():
    os.makedirs(SRC)
    shutil.copytree(TOP, SRC, dirs_exist_ok=True,
                    ignore=shutil.ignore_patterns("target", ".git", "debian", "vendor"))
    series = open(os.path.join(TOP, "debian/patches/series")).read().split()
    for p in series:
        r = run(["patch", "-p1", "-d", SRC, "-i", os.path.join(TOP, "debian/patches", p)])
        if r.returncode:
            sys.exit("patch %s failed:\n%s" % (p, r.stdout + r.stderr))
    open(os.path.join(SRC, "rust/Cargo.lock"), "a").close()


def vendor_everything():
    r = run(["cargo", "vendor", "--versioned-dirs", "--manifest-path",
             os.path.join(SRC, "rust/Cargo.toml"), FULL], env=ENV)
    if r.returncode:
        sys.exit("cargo vendor failed:\n" + r.stderr)
    return sorted(os.listdir(FULL))


def archive_packages():
    """{package: [versions]} for every librust-*-dev in the archive."""
    pk = collections.defaultdict(set)
    cur = None
    for line in run(["apt-cache", "dump"]).stdout.splitlines():
        if line.startswith("Package: "):
            cur = line[9:] if re.fullmatch(r"librust-.*-dev", line[9:]) else None
        elif line.startswith(" Version: ") and cur:
            pk[cur].add(line[10:])
    return pk


def add_archive_crate(pk, name):
    """Unpack every archive version of a crate. Returns False if it has none."""
    base = "librust-" + name.replace("_", "-").lower()
    want = ["%s=%s" % (p, v) for p in pk
            if re.fullmatch(re.escape(base) + r"(-\d[\d.]*)?-dev", p) for v in pk[p]]
    if not want:
        return False
    os.makedirs(DEBS, exist_ok=True)
    before = set(os.listdir(DEBS))
    run(["apt-get", "download"] + want, cwd=DEBS)
    for f in set(os.listdir(DEBS)) - before:
        run(["dpkg", "-x", os.path.join(DEBS, f), ARCH])
    return True


def build_merged(vendored):
    """Archive crates, except where a vendored copy of the same crate exists."""
    shutil.rmtree(MERGED, ignore_errors=True)
    os.makedirs(MERGED)
    reg = os.path.join(ARCH, REGISTRY)
    for d in os.listdir(reg) if os.path.isdir(reg) else []:
        if not any(d == v or d.startswith((v + "-", v + "+")) for v in vendored):
            os.symlink(os.path.join(reg, d), os.path.join(MERGED, d))
    for d in vendored:
        os.symlink(os.path.join(FULL, d), os.path.join(MERGED, d))
    os.makedirs(ENV["CARGO_HOME"], exist_ok=True)
    with open(os.path.join(ENV["CARGO_HOME"], "config.toml"), "w") as f:
        f.write('[source.crates-io]\nreplace-with = "m"\n[source.m]\ndirectory = "%s"\n' % MERGED)


def metadata():
    return run(["cargo", "metadata", "--offline", "--format-version", "1",
                "--no-default-features", "--manifest-path",
                os.path.join(SRC, "rust/Cargo.toml")], env=ENV)


def resolves(vendored):
    build_merged(vendored)
    open(os.path.join(SRC, "rust/Cargo.lock"), "w").close()
    return metadata().returncode == 0


def crate_of(d):
    return re.fullmatch(r"(.+)-\d+\.\d+\.\d+.*", d).group(1)


def dirs_of(all_dirs, names):
    return {d for d in all_dirs if crate_of(d) in names}


def initial_set(pk, all_dirs):
    """Resolve, vendoring a whole crate whenever the archive cannot serve it."""
    known = {crate_of(d) for d in all_dirs}
    for n in sorted(known):
        add_archive_crate(pk, n)
    forced, fetched = set(), set(known)
    for _ in range(200):
        build_merged(dirs_of(all_dirs, forced))
        open(os.path.join(SRC, "rust/Cargo.lock"), "w").close()
        r = metadata()
        if r.returncode == 0:
            return dirs_of(all_dirs, forced)
        err = r.stderr
        fm = re.search(r"depends on `([^`]+)` with feature `[^`]+` but `[^`]+` does not have", err)
        m = re.search(r"requirement `([A-Za-z0-9_-]+) =", err) or \
            re.search(r"no matching package named `([^`]+)`", err)
        name = fm.group(1) if fm else (m.group(1) if m else None)
        if name is None:
            sys.exit("cannot resolve:\n" + err[:800])
        if not fm and name not in fetched:
            # A crate only the archive's own versions need: fetch it.
            fetched.add(name)
            if add_archive_crate(pk, name):
                continue
        if name not in known:
            sys.exit("%s is needed, is not in the archive and is not in the lock file" % name)
        if name in forced:
            sys.exit("vendoring %s did not resolve:\n%s" % (name, err[:800]))
        forced.add(name)
    sys.exit("resolution did not converge")


def minimise(start, all_dirs):
    keep = set(start)
    # Generic crates first; the Matrix stack is what the vendoring is for.
    for d in sorted(start, key=lambda x: x.startswith(("matrix", "ruma"))):
        if resolves(keep - {d}):
            keep.discard(d)
    return keep


def closure_check(vendored):
    """The archive crates the build uses must all arrive through Build-Depends."""
    build_merged(vendored)
    open(os.path.join(SRC, "rust/Cargo.lock"), "w").close()
    meta = json.loads(metadata().stdout)
    nodes = {n["id"]: n for n in meta["resolve"]["nodes"]}
    pkgs = {p["id"]: p for p in meta["packages"]}
    reg = os.path.realpath(os.path.join(ARCH, REGISTRY))

    def dirof(p):
        return os.path.realpath(os.path.dirname(p["manifest_path"]))

    arch = {i for i, p in pkgs.items() if p["source"] and dirof(p).startswith(reg)}
    owner = {}
    for f in sorted(os.listdir(DEBS)):
        listing = run(["dpkg-deb", "-c", os.path.join(DEBS, f)]).stdout
        for line in listing.splitlines():
            m = re.search(REGISTRY + r"/([^/]+)/$", line)
            if m:
                owner[m.group(1)] = f.split("_")[0]
    parents = {i for i, p in pkgs.items() if p["source"] and i not in arch} | {meta["resolve"]["root"]}
    direct = {owner[os.path.basename(dirof(pkgs[d["pkg"]]))]
              for pid in parents for d in nodes[pid]["deps"] if d["pkg"] in arch}
    sim = run(["apt-get", "-s", "install", "--no-install-recommends"] + sorted(direct))
    if sim.returncode:
        sys.exit("apt cannot install the Build-Depends set:\n" + sim.stderr[:600])
    installed = set(re.findall(r"^Inst (\S+)", sim.stdout, re.M))
    used = {owner[os.path.basename(dirof(pkgs[i]))] for i in arch}
    uncovered = sorted(p for p in used - installed if run(["dpkg", "-s", p]).returncode)
    if uncovered:
        sys.exit("not covered by the Build-Depends closure: %s" % uncovered)
    return sorted(direct), len(arch)


def main():
    prepare_source()
    all_dirs = vendor_everything()
    pk = archive_packages()
    start = initial_set(pk, all_dirs)
    final = minimise(start, all_dirs)
    if not resolves(final):
        sys.exit("final set does not resolve")
    depends, n_arch = closure_check(final)
    with open(os.path.join(TOP, "debian/vendored-crates.txt"), "w") as f:
        f.write("\n".join(sorted(final)) + "\n")
    with open(os.path.join(TOP, "debian/cargo-build-depends.txt"), "w") as f:
        f.write("".join(" %s,\n" % d for d in depends))
    print("vendored %d crates, %d archive crates in %d Build-Depends; scratch: %s"
          % (len(final), n_arch, len(depends), WORK))


if __name__ == "__main__":
    main()
