#!/usr/bin/env python3
"""Merge all extracted .proto files into a single unified proto3 file,
resolving duplicate names by prefixing with the source module."""
import os, re, sys

PROTO_DIR = sys.argv[1]
OUT = sys.argv[2]

# Collect all message/enum blocks
all_blocks = []  # (kind, name, lines, source_file)
seen_names = {}
duplicates = set()

for fn in sorted(os.listdir(PROTO_DIR)):
    if not fn.endswith('.proto'):
        continue
    module = fn.replace('.proto','')
    txt = open(os.path.join(PROTO_DIR, fn)).read()
    lines = txt.split('\n')
    i = 0
    while i < len(lines):
        line = lines[i]
        # Skip syntax/package/import/option lines
        if re.match(r'^(syntax|package|import|option)\s', line):
            i += 1
            continue
        # Find message or enum blocks
        m = re.match(r'^(message|enum)\s+(\w+)\s*\{', line)
        if m:
            kind = m.group(1)
            name = m.group(2)
            depth = 1
            block_lines = [line]
            i += 1
            while i < len(lines) and depth > 0:
                depth += lines[i].count('{') - lines[i].count('}')
                block_lines.append(lines[i])
                i += 1
            # Handle duplicate names
            final_name = name
            if name in seen_names:
                final_name = module.capitalize() + name
                duplicates.add(name)
            seen_names[final_name] = True
            all_blocks.append((kind, final_name, block_lines, module))
        else:
            i += 1

# Write unified proto
with open(OUT, 'w') as f:
    f.write('syntax = "proto3";\n')
    f.write('option go_package = "github.com/antitbone/dwarfctl/proto;pb";\n\n')
    for kind, name, block_lines, module in all_blocks:
        f.write(f'// source: {module}\n')
        # Replace the first line's name if it was renamed
        block_lines[0] = re.sub(
            rf'^({kind})\s+\w+\s*\{{',
            rf'\1 {name} {{',
            block_lines[0])
        f.write('\n'.join(block_lines) + '\n\n')

print(f'Merged {len(all_blocks)} blocks into {OUT}')
print(f'Renamed duplicates: {duplicates}')