"""scripts/check_redistributable.py: what may reach the repository and each release variant."""
import os
from pathlib import Path
import subprocess
import sys
import tempfile
import unittest
import zipfile

ROOT = Path(__file__).resolve().parents[1]
CHECK = ROOT / 'scripts/check_redistributable.py'
# The version-resource layout of an NVIDIA DLL: UTF-16 key, padding, UTF-16 value.
NVIDIA_PE = (b'MZ' + b'\0' * 62 + 'CompanyName'.encode('utf-16-le') + b'\0\0'
             + 'NVIDIA Corporation'.encode('utf-16-le'))


def run(*args, cwd=None):
    return subprocess.run([sys.executable, str(CHECK), *args], cwd=cwd, text=True, capture_output=True)


def write(path, data=b'fixture'):
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_bytes(data)


class TreeTests(unittest.TestCase):
    def setUp(self):
        self.temporary = tempfile.TemporaryDirectory()
        self.repo = Path(self.temporary.name)
        env = {k: v for k, v in os.environ.items() if not k.startswith('GIT_')}
        self.git = lambda *a: subprocess.run(['git', '-C', str(self.repo), *a], check=True, env=env,
                                             capture_output=True)
        self.git('init', '-q')
        write(self.repo / 'engine/runtime.cpp', b'int main() {}\n')
        write(self.repo / '.gitignore', b'/build/\n')
        self.git('add', '-A')

    def tearDown(self):
        self.temporary.cleanup()

    def test_ignored_files_pass_and_unignored_model_data_fails(self):
        write(self.repo / 'build/k-model/weights.bin')
        self.assertEqual(run('tree', str(self.repo)).returncode, 0)
        write(self.repo / 'engine/k-model/weights.bin')
        result = run('tree', str(self.repo))
        self.assertEqual(result.returncode, 1)
        self.assertIn('engine/k-model/weights.bin: engine model data', result.stderr)

    def test_renamed_nvidia_binary_and_translation_are_found(self):
        write(self.repo / 'tools/helper.bin', NVIDIA_PE)
        write(self.repo / 'kernels/k/out.hip',
              b'// Generated by kernels/native/ptx2hip.py from NVIDIA PTX (k); 0 roundf idioms.\n')
        result = run('tree', str(self.repo))
        self.assertEqual(result.returncode, 1)
        self.assertIn('tools/helper.bin: Windows binary published by NVIDIA', result.stderr)
        self.assertIn('kernels/k/out.hip: translation of NVIDIA PTX', result.stderr)

    def test_staged_mode_checks_the_index_not_the_work_tree(self):
        write(self.repo / 'tools/helper.dll', NVIDIA_PE)
        self.git('add', 'tools/helper.dll')
        write(self.repo / 'tools/helper.dll', b'replaced after staging')
        result = run('tree', '--staged', str(self.repo))
        self.assertEqual(result.returncode, 1)
        self.assertIn('tools/helper.dll: Windows binary published by NVIDIA', result.stderr)
        self.git('rm', '-q', '-f', '--cached', 'tools/helper.dll')
        self.assertEqual(run('tree', '--staged', str(self.repo)).returncode, 0)


class ReleaseTests(unittest.TestCase):
    def stage(self, root, extra):
        write(root / 'd4r/nvngx.dll')  # d4r's shim under NVIDIA's name
        write(root / 'd4r/kernels/gfx1101/dltss_pwin_enc1_layer.hsaco')
        write(root / 'd4r/kernels/gfx1101/rrlite_enc1_4x4.hsaco')
        write(root / 'd4r/engine/k/input_k.spv')
        for name, data in extra.items():
            write(root / name, data)

    def check(self, variant, extra, archive=False):
        with tempfile.TemporaryDirectory() as temporary:
            root = Path(temporary) / 'stage'
            self.stage(root, extra)
            target = root
            if archive:
                target = Path(temporary) / 'release.zip'
                with zipfile.ZipFile(target, 'w') as output:
                    for file in root.rglob('*'):
                        output.write(file, file.relative_to(root).as_posix())
            return run('release', variant, str(target))

    def test_own_files_pass_every_variant(self):
        for variant in ('clean', 'installer', 'full'):
            for archive in (False, True):
                with self.subTest(variant=variant, archive=archive):
                    result = self.check(variant, {}, archive)
                    self.assertEqual(result.returncode, 0, result.stderr)

    def test_each_variant_holds_only_its_nvidia_files(self):
        dlls = {'d4r/nvngx_dlss.dll': NVIDIA_PE, 'd4r/ngx/_nvngx.dll': NVIDIA_PE}
        kernels = {'d4r/kernels/gfx1101/rrlite_post_3_1_mvlo_hdr.hsaco': b'code',
                   'd4r/kernels/gfx1101/hiluma_engine_input_depthinv_mvlo_hdr_v2_rel.hsaco': b'code'}
        for archive in (False, True):
            with self.subTest(archive=archive):
                clean = self.check('clean', {**dlls, **kernels}, archive)
                self.assertEqual(clean.returncode, 1)
                for name in {**dlls, **kernels}:
                    self.assertIn(name + ':', clean.stderr)
                self.assertEqual(self.check('installer', kernels, archive).returncode, 0)
                installer = self.check('installer', dlls, archive)
                self.assertEqual(installer.returncode, 1)
                for name in dlls:
                    self.assertIn(name + ':', installer.stderr)
                self.assertEqual(self.check('full', {**dlls, **kernels}, archive).returncode, 0)

    def test_nvidia_dll_elsewhere_and_dlss_values_fail_every_variant(self):
        refused = {'d4r/extra/nvngx_dlss.dll': NVIDIA_PE, 'd4r/engine/k/weights.bin': b'w',
                   'd4r/engine/k/model.json': b'{}', 'd4r/kernels/gfx1101/layer.ptx': b'.version 8.7'}
        for variant in ('clean', 'installer', 'full'):
            with self.subTest(variant=variant):
                result = self.check(variant, refused)
                self.assertEqual(result.returncode, 1)
                for name in refused:
                    self.assertIn(name + ':', result.stderr)


if __name__ == '__main__':
    unittest.main()