"""Read a named embedded FileDescriptorProto without executing its carrier."""
import argparse
import hashlib
import json
from pathlib import Path
from google.protobuf import descriptor_pb2, json_format


def varint(data, pos):
    value = 0
    for shift in range(0, 70, 7):
        byte = data[pos]
        pos += 1
        value |= (byte & 127) << shift
        if byte < 128:
            return value, pos
    raise ValueError('Oversized varint')


def extract(data, name):
    encoded = name.encode()
    occurrences = []
    pos = 0
    while (pos := data.find(encoded, pos)) >= 0:
        if len(encoded) < 128 and data[pos-2:pos] == bytes([10, len(encoded)]):
            occurrences.append(pos-2)
        pos += len(encoded)
    valid = []
    for start in occurrences:
        pos = start
        while pos < min(len(data), start+16*1024*1024):
            tag, pos = varint(data, pos)
            if tag & 7 != 2 or tag >> 3 not in (1, 2, 3, 4, 5, 6, 7, 8, 9, 12, 14):
                break
            size, pos = varint(data, pos)
            end = pos+size
            if end > len(data):
                break
            if tag >> 3 == 12 and data[pos:end] in (b'proto2', b'proto3'):
                desc = descriptor_pb2.FileDescriptorProto.FromString(data[start:end])
                if desc.name == name:
                    valid.append((start, end, desc))
                break
            pos = end
    if len(valid) != 1:
        raise ValueError(f'Expected one complete descriptor {name}, found {len(valid)}')
    return valid[0]


def main():
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument('--binary', type=Path, required=True)
    p.add_argument('--descriptor', required=True)
    p.add_argument('--output', type=Path, required=True)
    args = p.parse_args()
    data = args.binary.read_bytes()
    start, end, desc = extract(data, args.descriptor)
    args.output.mkdir(parents=True, exist_ok=True)
    stem = Path(args.descriptor).stem
    provenance = {'artifact_sha256':hashlib.sha256(data).hexdigest(), 'descriptor':desc.name, 'offset':start, 'length':end-start, 'descriptor_sha256':hashlib.sha256(data[start:end]).hexdigest()}
    (args.output/(stem+'.descriptor.pb')).write_bytes(data[start:end])
    (args.output/(stem+'.descriptor.json')).write_text(json.dumps({'provenance':provenance, 'descriptor':json_format.MessageToDict(desc, preserving_proto_field_name=True)}, indent=2), encoding='utf8')
    cmds = {v.name:v.number for enum in desc.enum_type if enum.name.startswith('Cmd') for v in enum.value}
    if cmds:
        (args.output/'cmdids.json').write_text(json.dumps({'provenance':provenance, 'commands':cmds},indent=2), encoding='utf8')
    print(json.dumps({**provenance, 'enums':len(desc.enum_type), 'messages':len(desc.message_type), 'commands':len(cmds)}))


if __name__ == '__main__':
    main()