#!/usr/bin/env python3
"""
Create a tar archive of node_modules files belonging to the given packages'
transitive dependency trees, written to stdout.  If no packages are given,
defaults to all entries in package.json dependencies.
"""

import argparse
import json
import logging
import os
import sys
import tarfile
from collections.abc import Iterator, Mapping
from pathlib import Path
from typing import Any

logger = logging.getLogger(__name__)


class PackageDB:
    def __init__(self) -> None:
        with Path('package-lock.json').open() as fp:
            self.packages: Mapping[str, Any] = json.load(fp)['packages']

    def depends(self) -> Iterator[str]:
        with Path('package.json').open() as fp:
            yield from json.load(fp).get('dependencies', {})

    def resolve(self, name: str, scope: Path) -> tuple[Path, Any]:
        while True:
            key = scope / 'node_modules' / name
            try:
                return key, self.packages[str(key)]
            except KeyError:
                if scope == Path('.'):
                    raise
                scope = scope.parent.parent  # skip past node_modules/pkg

    def files(self, name: str, scope: Path = Path('.')) -> Iterator[Path]:
        key, pkg = self.resolve(name, scope)

        for dirpath, dirnames, filenames in os.walk(key):
            # Don't descend into nested node_modules — those are separate lockfile entries.
            dirnames[:] = [d for d in dirnames if d != 'node_modules']
            for fname in filenames:
                yield Path(dirpath) / fname

        for dep_name in pkg.get('dependencies', {}):
            yield from self.files(dep_name, scope=key)


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('-d', '--debug', action='store_true')
    parser.add_argument('-C', '--directory', type=Path, default=Path('.'))
    parser.add_argument('--package-lock', action='store_true')
    parser.add_argument('packages', nargs='*', metavar='PACKAGE')
    args = parser.parse_args()

    logging.basicConfig(level=logging.DEBUG if args.debug else logging.WARNING)

    if sys.stdout.isatty():
        sys.exit('refusing to write tar to a terminal')

    os.chdir(args.directory)

    db = PackageDB()
    packages = args.packages or db.depends()

    files: set[Path] = set()
    if args.package_lock:
        files.add(Path('node_modules') / '.package-lock.json')
    for pkg in packages:
        files.update(db.files(pkg))

    with tarfile.open(mode='w|', fileobj=sys.stdout.buffer, format=tarfile.USTAR_FORMAT) as tf:
        for f in sorted(files):
            with f.open('rb') as fh:
                info = tf.gettarinfo(arcname=str(f), fileobj=fh)
                info.uid = info.gid = 0
                info.uname = info.gname = 'root'
                tf.addfile(info, fh)


if __name__ == '__main__':
    main()