#!/usr/bin/env python3
"""
Extract serialized FileDescriptorProto bytes from jadx-generated *Proto.java
files and regenerate clean .proto definitions using the real protobuf library.
"""
import re
import sys
import os
from google.protobuf import descriptor_pb2

# proto3 -> no "optional" label; proto2 prints labels.
TYPE_MAP = {
    1: 'double', 2: 'float', 3: 'int64', 4: 'uint64', 5: 'int32',
    6: 'fixed64', 7: 'fixed32', 8: 'bool', 9: 'string', 11: 'message',
    12: 'bytes', 13: 'uint32', 14: 'enum', 15: 'sfixed32', 16: 'sfixed64',
    17: 'sint32', 18: 'sint64',
}

def field_type_str(f):
    if f.type_name:
        return f.type_name.lstrip('.')
    return TYPE_MAP.get(f.type, f'scalar{f.type}')

def render_field(f, syntax):
    ty = field_type_str(f)
    label = ''
    if f.label == descriptor_pb2.FieldDescriptorProto.LABEL_REPEATED:
        if f.options.packed:
            label = 'packed repeated '
        else:
            label = 'repeated '
    elif syntax == 'proto2' and f.label == descriptor_pb2.FieldDescriptorProto.LABEL_REQUIRED:
        label = 'required '
    elif syntax == 'proto2':
        label = 'optional '
    return f'  {label}{ty} {f.name} = {f.number};'

def render_enum(e, indent):
    lines = [f'{indent}enum {e.name} {{']
    for v in e.value:
        lines.append(f'{indent}  {v.name} = {v.number};')
    lines.append(f'{indent}}}')
    return '\n'.join(lines)

def render_message(m, indent, syntax):
    lines = [f'{indent}message {m.name} {{']
    for oneof in m.oneof_decl:
        lines.append(f'{indent}  // oneof {oneof.name}')
    for nested in m.enum_type:
        lines.append(render_enum(nested, indent + '  '))
    for nested in m.nested_type:
        lines.append(render_message(nested, indent + '  ', syntax))
    for f in m.field:
        lines.append(render_field(f, syntax))
    lines.append(f'{indent}}}')
    return '\n'.join(lines)

def render_file(fd):
    out = [f'// source: {fd.name}']
    if fd.package:
        out.append(f'package {fd.package};')
    syntax = fd.syntax or 'proto3'
    out.append(f'syntax = "{syntax}";')
    if fd.dependency:
        out.append('')
        for dep in fd.dependency:
            out.append(f'import "{dep}";')
    out.append('')
    for e in fd.enum_type:
        out.append(render_enum(e, ''))
        out.append('')
    for m in fd.message_type:
        out.append(render_message(m, '', syntax))
        out.append('')
    return '\n'.join(out)

DESC_RE = re.compile(
    r'internalBuildGeneratedFileFrom\(new String\[\]\{(.*?)\}\s*,\s*new Descriptors',
    re.DOTALL)

def decode_java_string_literal_body(lit):
    """Decode a Java string-literal body (between quotes) to raw bytes."""
    out = bytearray()
    i = 0; n = len(lit)
    while i < n:
        c = lit[i]
        if c == '\\':
            nxt = lit[i+1]
            simple = {'n': 10, 'r': 13, 't': 9, '"': 34, "'": 39, '\\': 92, '0': 0, 'b': 8, 'f': 12}
            if nxt in simple:
                out.append(simple[nxt]); i += 2
            elif nxt == 'u':
                out += chr(int(lit[i+2:i+6], 16)).encode('latin-1', 'replace')
                i += 6
            else:
                out.append(ord(nxt) & 0xff); i += 2
        else:
            out += c.encode('latin-1', 'replace'); i += 1
    return bytes(out)

def extract_descriptor_bytes(java_path):
    txt = open(java_path, encoding='utf-8', errors='replace').read()
    m = DESC_RE.search(txt)
    if not m:
        return None
    body = m.group(1)
    lits = re.findall(r'"((?:[^"\\]|\\.)*)"', body)
    # protobuf concatenates ALL adjacent string elements into ONE descriptor
    combined = b''.join(decode_java_string_literal_body(l) for l in lits)
    return combined if combined else None

def main():
    proto_dir = sys.argv[1]
    out_dir = sys.argv[2]
    os.makedirs(out_dir, exist_ok=True)
    total_msgs = 0
    for fn in sorted(os.listdir(proto_dir)):
        if not fn.endswith('Proto.java'):
            continue
        path = os.path.join(proto_dir, fn)
        raw = extract_descriptor_bytes(path)
        base = fn[:-5].replace('Proto', '')
        if not raw:
            print(f'  [WARN] no descriptor in {fn}')
            continue
        try:
            fd_set = descriptor_pb2.FileDescriptorSet()
            fd = fd_set.file.add()
            fd.MergeFromString(raw)
        except Exception as ex:
            # may be a FileDescriptorSet directly, or multiple
            try:
                fd_set.ParseFromString(raw)
                if not fd_set.file:
                    raise
            except Exception as ex2:
                print(f'  [ERR ] {fn}: {ex2}')
                continue
        out_path = os.path.join(out_dir, base + '.proto')
        chunks = []
        for one_fd in fd_set.file:
            chunks.append(render_file(one_fd))
            total_msgs += len(one_fd.message_type)
        with open(out_path, 'w', encoding='utf-8') as f:
            f.write('\n\n'.join(chunks) + '\n')
        print(f'  [OK] {fn} -> {base}.proto ({len(raw)} bytes, {len(fd_set.file)} file(s), {sum(len(x.message_type) for x in fd_set.file)} msgs)')
    print(f'\nTotal messages recovered: {total_msgs}')

if __name__ == '__main__':
    main()