"""Content-addressed Runtime source and complete installed dependency manifests."""
import argparse
import hashlib
import json
import struct
from pathlib import Path


def sha(path):
    return hashlib.sha256(path.read_bytes()).hexdigest()


def platform(root):
    if (root / 'lib/yolo_runtime.dll').exists():
        return 'windows-x64'
    binary = (root / 'lib/libyolo_runtime.so').read_bytes()
    if binary[:4] != b'\x7fELF':
        raise ValueError('Installed Runtime is not an ELF library')
    machine = struct.unpack_from('<H', binary, 18)[0]
    if machine not in (183, 62):
        raise ValueError('Unsupported installed Runtime architecture')
    return 'linux-arm64' if machine == 183 else 'linux-x64'


def source(root):
    files = [p for part in ('include', 'src', 'backends', 'adapters', 'cmake', 'scripts')
             for p in (root / part).rglob('*') if p.is_file() and '__pycache__' not in p.parts]
    files += [root / 'CMakeLists.txt', *root.glob('*.def')]
    hashes = {p.relative_to(root).as_posix(): sha(p) for p in sorted(files)}
    revision = hashlib.sha256(json.dumps(hashes, sort_keys=True).encode()).hexdigest()
    return dict(schema_version=1, version='5.0.0', abi=5, source_revision=revision, files=hashes)


def verify(root, expected=None):
    data = json.loads((root / 'runtime-manifest.json').read_text(encoding='utf-8'))
    if data['abi'] != 5 or (expected and data['source_revision'] != expected):
        raise ValueError('Runtime ABI/source revision mismatch')
    for name, digest in data['files'].items():
        target = (root / name).resolve()
        if root.resolve() not in target.parents or not target.is_file() or sha(target) != digest:
            raise ValueError('Runtime missing or changed dependency: ' + name)
    required = ['include/yolo_runtime.h', 'include/yolo_runtime_types.h', 'source-identity.json']
    required += (['lib/yolo_runtime.dll', 'lib/yolo_backend_onnx_cpu.dll', 'lib/onnxruntime.dll']
                 if data['platform'] == 'windows-x64' else ['lib/libyolo_runtime.so'])
    if any(name not in data['files'] for name in required):
        raise ValueError('Runtime manifest omits a required dependency')
    print('RUNTIME_MANIFEST_OK', data['source_revision'])


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('mode', choices=['source', 'write', 'verify'])
    parser.add_argument('--root', type=Path, required=True)
    parser.add_argument('--output', type=Path)
    parser.add_argument('--expected')
    args = parser.parse_args()
    root = args.root.resolve()
    if args.mode == 'source':
        args.output.write_text(json.dumps(source(root), indent=2), encoding='utf-8')
    elif args.mode == 'write':
        identity = json.loads((root / 'source-identity.json').read_text(encoding='utf-8'))
        files = {p.relative_to(root).as_posix(): sha(p) for p in sorted(root.rglob('*'))
                 if p.is_file() and p.name != 'runtime-manifest.json'}
        data = dict(schema_version=1, version=identity['version'], abi=identity['abi'],
                    source_revision=identity['source_revision'],
                    platform=platform(root), files=files)
        (root / 'runtime-manifest.json').write_text(json.dumps(data, indent=2), encoding='utf-8')
        verify(root)
    else:
        verify(root, args.expected)


if __name__ == '__main__':
    main()