#!/usr/bin/env python3
# Copyright 2023 The Fuchsia Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""Rewrite protobuf files to have monotonically increasing field numbers.

This is safe to do for recipe property protobuf files because they are only ever
used to de/serialize JSON-encoding protos, which are only sensitive to field
names, not field numbers.
"""

from __future__ import annotations

import argparse
import re

_ONEOF = object()


def renumber_fields(lines):
    # Copy the lines for modification.
    newlines = list(lines)
    # levels is a stack that tracks the next field number to use within a given
    # scope.
    levels = []
    for i, line in enumerate(lines):
        stripped = line.split("//")[0].strip()
        if not stripped:
            continue
        if stripped.endswith("{"):
            if stripped.startswith("enum"):
                # First field number for enums is 0.
                levels.append(0)
            elif stripped.startswith("oneof"):
                # oneofs share numbering with the parent level
                levels.append(_ONEOF)
            elif stripped.startswith("message"):
                # First field number for messages is 1.
                levels.append(1)
            elif stripped.startswith("option"):
                # option clauses shouldn't contain any numbered fields.
                levels.append(None)
            else:
                raise ValueError(f"invalid line: {line}")
            continue
        if stripped.startswith("}"):
            levels.pop()
            continue

        # A single line may contain multiple fields.
        for match in re.finditer(
            r"((?P<reserved>reserved\s+\w+(\s*(,|to\s)\s*\w+)*)|(?P<field>\w+ =) \d+(?P<option>\s*(\[.*\])?));",
            stripped,
        ):
            if match.group("reserved"):
                # Reserved fields are not necessary in recipe property protos.
                newlines[i] = line.replace(match.group(0), "")
                continue

            idx = 0
            for j, level in enumerate(levels):
                if level != _ONEOF:
                    idx = j
            level = levels[idx]
            new = f"{match.group('field')} {level}{match.group('option') or ''};"
            newlines[i] = line.replace(match.group(0), new)
            levels[idx] += 1

    return newlines


def main():
    parser = argparse.ArgumentParser(description="Renumber .proto file fields.")
    parser.add_argument("filename", metavar="FILE", type=str, help="The file to modify")
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Print result to stdout instead of writing to the file",
    )
    args = parser.parse_args()

    with open(args.filename, encoding="utf-8") as f:
        lines = f.readlines()

    new_lines = renumber_fields(lines)

    if args.dry_run:
        for l in new_lines:
            print(l, end="")
    else:
        with open(args.filename, "w", encoding="utf-8") as f:
            f.writelines(new_lines)


if __name__ == "__main__":
    main()