"""Verify the locked ARM64 wheelhouse before offline delivery."""

import argparse
from email.parser import BytesParser
from email.policy import default
from hashlib import sha256
from pathlib import Path
import tomllib
from zipfile import ZipFile

from packaging.markers import default_environment
from packaging.requirements import Requirement
from packaging.tags import compatible_tags, cpython_tags
from packaging.utils import canonicalize_name, parse_wheel_filename


ROOT = Path(__file__).resolve().parents[1]
PLATFORMS = [f'manylinux_2_{version}_aarch64' for version in range(28, 16, -1)]
PLATFORMS += ['manylinux2014_aarch64', 'linux_aarch64']
TAGS = set(cpython_tags(python_version=(3, 14), abis=['cp314'], platforms=PLATFORMS))
TAGS.update(compatible_tags(python_version=(3, 14), interpreter='cp314', platforms=PLATFORMS))


def verify(requirements_path, wheelhouse, lock_path):
    required = {}
    for line in requirements_path.read_text(encoding='utf-8').splitlines():
        if not line or line.startswith('#'):
            continue
        requirement = Requirement(line)
        name = canonicalize_name(requirement.name)
        if name in required or len(requirement.specifier) != 1:
            raise ValueError(f'invalid_target_requirement:{name}')
        required[name] = requirement.specifier

    with lock_path.open('rb') as stream:
        locked = tomllib.load(stream)['package']
    hashes = {}
    for package in locked:
        key = (canonicalize_name(package['name']), package['version'])
        hashes.setdefault(key, set()).update(
            wheel['hash'].removeprefix('sha256:') for wheel in package.get('wheels', [])
        )

    seen = set()
    total = 0
    installed = {}
    dependencies = []
    for path in sorted(wheelhouse.iterdir()):
        if not path.is_file() or path.suffix != '.whl':
            raise ValueError(f'unexpected_wheelhouse_file:{path.name}')
        name, version, _, tags = parse_wheel_filename(path.name)
        if name not in required or version not in required[name] or name in seen:
            raise ValueError(f'unexpected_wheel:{path.name}')
        if not tags & TAGS:
            raise ValueError(f'incompatible_wheel:{path.name}')
        digest = sha256()
        with path.open('rb') as stream:
            for block in iter(lambda: stream.read(1024 * 1024), b''):
                digest.update(block)
        if digest.hexdigest() not in hashes.get((name, str(version)), set()):
            raise ValueError(f'lock_hash_mismatch:{path.name}')
        with ZipFile(path) as archive:
            metadata_names = [entry for entry in archive.namelist()
                              if entry.count('/') == 1 and entry.endswith('.dist-info/METADATA')]
            if len(metadata_names) != 1:
                raise ValueError(f'invalid_wheel_metadata:{path.name}')
            metadata = BytesParser(policy=default).parsebytes(archive.read(metadata_names[0]))
            dependencies.extend((name, Requirement(value)) for value in metadata.get_all('Requires-Dist', []))
        seen.add(name)
        installed[name] = version
        total += path.stat().st_size
    if seen != set(required):
        raise ValueError(f'missing_wheels:{sorted(set(required) - seen)}')
    environment = default_environment()
    environment.update({
        'extra': '', 'implementation_name': 'cpython', 'os_name': 'posix',
        'platform_machine': 'aarch64', 'platform_python_implementation': 'CPython',
        'platform_system': 'Linux', 'python_full_version': '3.14.5',
        'python_version': '3.14', 'sys_platform': 'linux',
    })
    for parent, requirement in dependencies:
        if requirement.marker and not requirement.marker.evaluate(environment):
            continue
        dependency = canonicalize_name(requirement.name)
        if dependency not in installed or installed[dependency] not in requirement.specifier:
            raise ValueError(f'missing_dependency:{parent}:{requirement}')
    return len(seen), total


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('requirements', type=Path)
    parser.add_argument('wheelhouse', type=Path)
    parser.add_argument('--lock', type=Path, default=ROOT / 'uv.lock')
    args = parser.parse_args()
    count, total = verify(args.requirements, args.wheelhouse, args.lock)
    print(f'Locked ARM64 wheelhouse verified: {count} wheels, {total} bytes.')


if __name__ == '__main__':
    main()