#!/usr/bin/env python3
"""Refuse NVIDIA property and data read from it where d4r publishes files.

d4r's repository must not contain NVIDIA's software, code compiled or translated from NVIDIA's PTX, or values
read from NVIDIA's DLSS; each release variant may contain only what scripts/package_release.sh says it contains
(see NOTICE). This checks a file set for the rest.

usage: check_redistributable.py tree [--staged] [REPO]              files git would commit (all, or the index)
       check_redistributable.py release clean|installer|full PATH  a staged release folder or its zip

The variants are those of scripts/package_release.sh's D4R_BUNDLE_NVIDIA: `clean` (0) holds nothing of NVIDIA's;
`installer` (the default) adds texture kernels built from NVIDIA's PTX, and install.sh downloads NVIDIA's DLLs;
`full` (1) also bundles the two DLLs. Values read from DLSS are refused in every variant.
Exit status 1 lists every finding; 0 means none were found. This is a safeguard against mistakes, not proof
that a file set is free of third-party rights.
"""
import fnmatch
import re
import subprocess
import sys
import zipfile
from pathlib import Path, PurePosixPath

# Files the native engine's tools write from NVIDIA's DLSS (engine/README.md); NOTICE names them.
MODEL_DATA = {'weights.bin', 'offsets.bin', 'lut.bin', 'hipnet.bin', 'model.json'}
# ... and the M / L and Ray Reconstruction models' (engine/compile_m.py, compile_rr.py, compile_rr2.py)
MODEL_DATA_PATTERNS = ('weights8.bin', 'weights16.bin', 'weights-*.bin', '*.args', 'enc0-aux.bin', '*.generated.*')
# Compiled or extracted NVIDIA GPU code.
NVIDIA_CODE_SUFFIXES = {'.ptx', '.cubin', '.fatbin', '.nvngx'}
# Kernels kernels/build.sh builds from NVIDIA's PTX (tex/l families and the PTX-translated hiluma kernels).
# kernels/build.sh k and m build only d4r's own layers (dltss_pwin_*, rrlite_enc1/enc2/enc3/dec1/dec2).
PTX_KERNELS = ('hiluma_engine_*', 'rrlite_enc0_*', 'rrlite_dec0_*', 'rrlite_post_*', 'rrlite_downsample_*')
# Headers of NVIDIA's NGX/DLSS SDK, which the probes in tools/ build against from the user's own copy.
SDK_HEADERS = ('nvsdk_ngx*.h',)
# The comment d4r's PTX translators put in the first lines of their output.
TRANSLATION_MARKER = re.compile(rb'^// Generated by \S+ from NVIDIA PTX', re.MULTILINE)
TRANSLATION_HEADER_LINES = 5
NVIDIA_DLL = re.compile(r'_?nvngx.*\.dll', re.IGNORECASE)
SHIM_RELEASE_PATH = 'd4r/nvngx.dll'
FULL_RELEASE_NVIDIA = {'d4r/nvngx_dlss.dll', 'd4r/ngx/_nvngx.dll'}
# Large binaries in the tree are not d4r's style (sources, patches, docs); one is most likely captured data.
TREE_BINARY_LIMIT = 1 << 20
COMPANY = 'CompanyName'.encode('utf-16-le')
NVIDIA_UTF16 = 'NVIDIA'.encode('utf-16-le')


def nvidia_pe(data):
    """A Windows binary whose version resource names NVIDIA as its company."""
    if not data.startswith(b'MZ'):
        return False
    at = data.find(COMPANY)
    while at >= 0:
        # VS_VERSIONINFO: the key, NUL padding to a 32-bit boundary, then the UTF-16 value.
        if NVIDIA_UTF16 in data[at + len(COMPANY):at + len(COMPANY) + 64]:
            return True
        at = data.find(COMPANY, at + 1)
    return False


def findings_for(path, read, release=None):
    """Problems with one file. `path` is POSIX and relative; `read()` returns its bytes."""
    name = PurePosixPath(path).name
    lower = name.lower()
    suffix = PurePosixPath(lower).suffix
    full = release == 'full'
    problems = []
    if lower in MODEL_DATA or any(fnmatch.fnmatch(lower, pattern) for pattern in MODEL_DATA_PATTERNS):
        problems.append('engine model data read from NVIDIA\'s DLSS')
    if suffix in NVIDIA_CODE_SUFFIXES:
        problems.append('NVIDIA GPU code')
    if any(fnmatch.fnmatch(lower, pattern) for pattern in SDK_HEADERS):
        problems.append('NVIDIA SDK header')
    if NVIDIA_DLL.fullmatch(name) and path != SHIM_RELEASE_PATH and not (full and path in FULL_RELEASE_NVIDIA):
        problems.append('NVIDIA NGX/DLSS library')
    if release is None and suffix in {'.hsaco', '.co', '.spv'}:
        problems.append('compiled code object (build output)')
    if release == 'clean' and any(fnmatch.fnmatch(lower, pattern) for pattern in PTX_KERNELS):
        problems.append('kernel built from NVIDIA\'s PTX')
    data = read()
    if TRANSLATION_MARKER.search(b'\n'.join(data[:4096].split(b'\n')[:TRANSLATION_HEADER_LINES])):
        problems.append('translation of NVIDIA PTX')
    if nvidia_pe(data) and not (full and path in FULL_RELEASE_NVIDIA):
        problems.append('Windows binary published by NVIDIA')
    if release is None and len(data) > TREE_BINARY_LIMIT and b'\0' in data[:8192]:
        problems.append(f'binary file over {TREE_BINARY_LIMIT >> 20} MiB')
    return problems


def git(repo, *args):
    return subprocess.run(['git', '-C', str(repo), *args], check=True, capture_output=True).stdout


def tree_files(repo, staged):
    if staged:
        names = git(repo, 'diff', '--cached', '--name-only', '--diff-filter=ACMR', '-z')
        reader = lambda path: lambda: git(repo, 'show', f':{path}')
    else:
        names = git(repo, 'ls-files', '-z', '--cached', '--others', '--exclude-standard')
        reader = lambda path: lambda: (repo / path).read_bytes()
    for path in names.decode().split('\0'):
        # Gitlinks (external/MinHook) and files deleted from the work tree have nothing to read.
        if path and (staged or (repo / path).is_file()):
            yield path, reader(path)


def release_files(root):
    if root.is_file():
        archive = zipfile.ZipFile(root)
        for info in archive.infolist():
            if not info.is_dir():
                yield info.filename, (lambda info=info: archive.read(info))
        return
    for file in sorted(root.rglob('*')):
        if file.is_file():
            yield file.relative_to(root).as_posix(), file.read_bytes


def main(argv):
    if len(argv) >= 1 and argv[0] == 'tree':
        staged = '--staged' in argv[1:]
        rest = [a for a in argv[1:] if a != '--staged']
        if len(rest) > 1:
            sys.exit(__doc__)
        repo = Path(rest[0] if rest else '.').resolve()
        repo = Path(git(repo, 'rev-parse', '--show-toplevel').decode().strip())
        files, release, where = tree_files(repo, staged), None, 'the index' if staged else str(repo)
    elif len(argv) == 3 and argv[0] == 'release' and argv[1] in ('clean', 'installer', 'full'):
        root = Path(argv[2])
        if not root.exists():
            sys.exit(f'{root} does not exist')
        files, release, where = release_files(root), argv[1], str(root)
    else:
        sys.exit(__doc__)
    found = 0
    for path, read in files:
        for problem in findings_for(path, read, release):
            print(f'{path}: {problem}', file=sys.stderr)
            found += 1
    if found:
        print(f'{found} file(s) in {where} must not be published (see NOTICE)', file=sys.stderr)
        return 1
    return 0


if __name__ == '__main__':
    sys.exit(main(sys.argv[1:]))