#!/usr/bin/env python3
"""Verify the exported model bundle using only the standard library, on CPU."""
import argparse
import hashlib
import json
from pathlib import Path
import sys


def verify(root):
    root = Path(root).resolve()
    manifest = json.loads((root / 'bundle_manifest.json').read_text())
    if manifest.get('status') != 'complete':
        raise ValueError('bundle is not marked complete')
    limit = manifest['max_file_bytes']
    total = 0
    for row in manifest['files']:
        relative = Path(row['path'])
        if relative.is_absolute() or '..' in relative.parts:
            raise ValueError('unsafe manifest path')
        path = root / relative
        if path.is_symlink() or not path.is_file() or path.stat().st_size != row['bytes']:
            raise ValueError('missing/changed file: ' + str(relative))
        if path.stat().st_size > limit:
            raise ValueError('file exceeds upload shard limit: ' + str(relative))
        digest = hashlib.sha256()
        with path.open('rb') as handle:
            for block in iter(lambda: handle.read(8*1024*1024), b''):
                digest.update(block)
        if digest.hexdigest() != row['sha256']:
            raise ValueError('SHA256 differs: ' + str(relative))
        total += row['bytes']
    print(json.dumps({'status': 'verified', 'files': len(manifest['files']),
                      'file_bytes': total, 'max_file_bytes': limit,
                      'GPU_used': False, 'model_loaded': False}, indent=2))
    return manifest


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--bundle', type=Path, default=Path(__file__).resolve().parents[1])
    args = parser.parse_args()
    verify(args.bundle)


if __name__ == '__main__':
    main()