"""Pack and verify the locked Kylin ARM64 CPU environment archive."""

import argparse
from hashlib import sha256
from pathlib import Path
import tarfile
import zipfile

from scripts.verify_offline_wheels import verify as verify_wheels


ROOT = Path(__file__).resolve().parents[1]
UV_ARCHIVE_SHA256 = 'bb66cb52e7b1823aed1183630d8d8e5c958840d584a4c55ec10a4cfc168dcca2'
UV_MEMBER = 'uv-aarch64-unknown-linux-gnu/uv'


def digest(path):
    result = sha256()
    with path.open('rb') as stream:
        for block in iter(lambda: stream.read(1024 * 1024), b''):
            result.update(block)
    return result.hexdigest()


def pack(source, output):
    source, output = Path(source), Path(output)
    if output.exists():
        raise ValueError('environment_archive_exists')
    requirements = source / 'requirements-arm64-cpu.txt'
    wheelhouse = source / 'wheelhouse'
    verify_wheels(requirements, wheelhouse, ROOT / 'uv.lock')
    uv_archive = source / 'uv-aarch64-unknown-linux-gnu.tar.gz'
    if digest(uv_archive) != UV_ARCHIVE_SHA256:
        raise ValueError('uv_archive_hash_mismatch')
    with tarfile.open(uv_archive, 'r:gz') as archive:
        member = archive.getmember(UV_MEMBER)
        if not member.isfile() or member.size > 100_000_000:
            raise ValueError('invalid_uv_binary')
        uv_bytes = archive.extractfile(member).read()
    if uv_bytes[:4] != b'\x7fELF' or int.from_bytes(uv_bytes[18:20], 'little') != 183:
        raise ValueError('uv_not_aarch64_elf')

    files = {'uv': uv_bytes}
    paths = {
        'requirements-arm64-cpu.txt': requirements,
        'install_offline_env.sh': ROOT / 'scripts/install_offline_env.sh',
    }
    paths.update({f'wheelhouse/{path.name}': path for path in sorted(wheelhouse.glob('*.whl'))})
    sums = {name: sha256(content).hexdigest() for name, content in files.items()}
    sums.update({name: digest(path) for name, path in paths.items()})
    checksums = ''.join(f'{value}  {name}\n' for name, value in sorted(sums.items()))

    output.parent.mkdir(parents=True, exist_ok=True)
    partial = output.with_name(output.name + '.partial')
    if partial.exists():
        raise ValueError('partial_environment_archive_exists')
    try:
        with zipfile.ZipFile(partial, 'w', allowZip64=True) as archive:
            archive.writestr('uv', uv_bytes, compress_type=zipfile.ZIP_STORED)
            for name, path in sorted(paths.items()):
                compression = zipfile.ZIP_STORED if path.suffix == '.whl' else zipfile.ZIP_DEFLATED
                archive.write(path, name, compress_type=compression)
            archive.writestr('SHA256SUMS', checksums)
        verify(partial)
        partial.replace(output)
    finally:
        partial.unlink(missing_ok=True)
    return len(sums)


def verify(path):
    with zipfile.ZipFile(path) as archive:
        lines = archive.read('SHA256SUMS').decode('ascii').splitlines()
        expected = {}
        for line in lines:
            checksum, name = line.split('  ', 1)
            if (len(checksum) != 64 or name in expected or name.startswith('/')
                    or '..' in Path(name).parts or '\\' in name):
                raise ValueError('invalid_environment_manifest')
            expected[name] = checksum
        if set(archive.namelist()) != set(expected) | {'SHA256SUMS'}:
            raise ValueError('unexpected_environment_members')
        for name, checksum in expected.items():
            with archive.open(name) as stream:
                actual = sha256()
                for block in iter(lambda: stream.read(1024 * 1024), b''):
                    actual.update(block)
            if actual.hexdigest() != checksum:
                raise ValueError(f'environment_member_hash_mismatch:{name}')
    return len(expected)


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    commands = parser.add_subparsers(dest='action', required=True)
    pack_parser = commands.add_parser('pack')
    pack_parser.add_argument('source', type=Path)
    pack_parser.add_argument('output', type=Path)
    verify_parser = commands.add_parser('verify')
    verify_parser.add_argument('archive', type=Path)
    args = parser.parse_args()
    count = pack(args.source, args.output) if args.action == 'pack' else verify(args.archive)
    print(f'Offline ARM64 environment archive verified: {count} files.')


if __name__ == '__main__':
    main()