#!/usr/bin/env python3
"""Inject nanopb .options constraints as inline field options into a .proto file.

The nanopb .options file format is specific to the nanopb C generator and is
ignored by standard protoc (including --python_out). By injecting the options
directly into the proto file's field declarations, protoc will embed them in
the serialized descriptor, making them accessible in Python via:

    from meshtastic.protobuf import mesh_pb2, nanopb_pb2
    field = mesh_pb2.DESCRIPTOR.message_types_by_name['User'].fields_by_name['long_name']
    opts = field.GetOptions().Extensions[nanopb_pb2.nanopb]
    print(opts.max_size)  # 40

Usage:
    inject_nanopb_options.py <options_file> <proto_file>

The proto_file is modified in-place. Intended to operate on temporary copies
generated by regen-protobufs.sh, not the source .proto files.
"""

import re
import sys
from pathlib import Path
from typing import Any, Dict, List, Tuple

# IntSize enum values from nanopb.proto
INT_SIZE_ENUM = {8: "IS_8", 16: "IS_16", 32: "IS_32", 64: "IS_64"}

# Options that are valid proto FieldOptions and useful outside of C code generation.
# We skip C-only options (anonymous_oneof, no_unions, skip_message, packed_struct,
# packed_enum, mangle_names, callback_datatype, callback_function, descriptorsize,
# type_override) since they either don't apply to proto fields or are C-specific.
FIELD_OPTIONS = frozenset(
    {
        "max_size",
        "max_length",
        "max_count",
        "int_size",
        "fixed_length",
        "fixed_count",
        "long_names",
        "proto3",
        "default_has",
        "sort_by_tag",
        "msgid",
    }
)


def parse_value(s: str) -> Any:
    """Convert an option value string to an appropriate Python type."""
    if re.fullmatch(r"-?[0-9]+", s):
        return int(s)
    if s.lower() == "true":
        return True
    if s.lower() == "false":
        return False
    return s


def parse_options_file(
    path: Path,
) -> Tuple[Dict[Tuple[str, ...], Dict[str, Any]], Dict[str, Dict[str, Any]]]:
    """Parse a nanopb .options file.

    Returns:
        specific: maps (message_path..., field_name) -> {option: value}
                  e.g. ('User', 'long_name') or ('Route', 'Link', 'uid')
        wildcard: maps field_name -> {option: value}
                  applies to any field with this name in the same proto file
    """
    specific: Dict[Tuple[str, ...], Dict[str, Any]] = {}
    wildcard: Dict[str, Dict[str, Any]] = {}

    with open(path, encoding="utf-8") as f:
        for line in f:
            # Strip inline comments
            comment_pos = line.find("#")
            if comment_pos >= 0:
                line = line[:comment_pos]
            line = line.strip().lstrip("*").strip()
            if not line:
                continue

            tokens = line.split()
            if len(tokens) < 2:
                continue

            pattern = tokens[0]
            opts: Dict[str, Any] = {}
            for tok in tokens[1:]:
                if ":" in tok:
                    k, v = tok.split(":", 1)
                    if k in FIELD_OPTIONS:
                        opts[k] = parse_value(v)

            if not opts:
                continue

            if "." in pattern:
                # e.g. "User.long_name" -> key=('User', 'long_name')
                # or "Route.Link.uid"   -> key=('Route', 'Link', 'uid')
                parts = tuple(pattern.split("."))
                if parts in specific:
                    specific[parts].update(opts)
                else:
                    specific[parts] = opts
            else:
                # wildcard: applies to any field with this name
                if pattern in wildcard:
                    wildcard[pattern].update(opts)
                else:
                    wildcard[pattern] = opts

    return specific, wildcard


def format_nanopb_opts(opts: Dict[str, Any]) -> str:
    """Format a dict of nanopb options as a proto field options annotation string."""
    parts = []
    for k, v in opts.items():
        if k == "int_size":
            enum_val = INT_SIZE_ENUM.get(v, f"IS_{v}")
            parts.append(f"(nanopb).int_size = {enum_val}")
        elif isinstance(v, bool):
            parts.append(f"(nanopb).{k} = {'true' if v else 'false'}")
        else:
            parts.append(f"(nanopb).{k} = {v}")
    return ", ".join(parts)


def message_path_matches(
    context_stack: List[Tuple[str, str]], msg_path: Tuple[str, ...]
) -> bool:
    """Check if the current message context ends with msg_path.

    context_stack entries are ('message'|'oneof'|'enum', name).
    msg_path is the tuple of message names from the options pattern,
    e.g. ('User',) or ('Route', 'Link').
    """
    msg_names = [name for kind, name in context_stack if kind == "message"]
    n = len(msg_path)
    return len(msg_names) >= n and tuple(msg_names[-n:]) == msg_path


def inject_into_proto(
    content: str,
    specific: Dict[Tuple[str, ...], Dict[str, Any]],
    wildcard: Dict[str, Dict[str, Any]],
    nanopb_import_path: str,
) -> str:
    """Inject nanopb field options into a proto file's text content.

    Adds an import for nanopb.proto if not already present.
    Returns the modified content.
    """
    if not specific and not wildcard:
        return content

    lines = content.split("\n")

    # Check if nanopb is already imported (after sed fixup, it will be
    # 'meshtastic/protobuf/nanopb.proto')
    nanopb_already_imported = any(
        "nanopb.proto" in line
        for line in lines
        if line.strip().startswith("import")
    )

    # Track the index of the last import line so we can insert after it
    last_import_idx = max(
        (
            i
            for i, line in enumerate(lines)
            if line.strip().startswith("import ") and line.strip().endswith(";")
        ),
        default=-1,
    )

    # State
    context_stack: List[Tuple[str, str]] = []  # ('message'|'oneof'|'enum', name)
    result: List[str] = []
    import_added = nanopb_already_imported

    # Patterns for proto structural elements
    message_re = re.compile(r"^(\s*)message\s+(\w+)\s*\{")
    oneof_re = re.compile(r"^(\s*)oneof\s+(\w+)\s*\{")
    enum_re = re.compile(r"^(\s*)enum\s+(\w+)\s*\{")
    close_re = re.compile(r"^\s*\}")

    # Pattern for field declarations:
    #   indent  [optional|repeated]  type  name  =  number  [options]  ;
    # We exclude map<> fields (different syntax, nanopb handles them differently)
    # and enum value lines (no type keyword before the name).
    field_re = re.compile(
        r"^(\s*)"  # (1) indent
        r"(optional\s+|repeated\s+)?"  # (2) optional qualifier
        r"([\w.]+)\s+"  # (3) field type (possibly qualified like google.protobuf.Any)
        r"(\w+)\s*"  # (4) field name
        r"=\s*(\d+)"  # (5) field number
        r"(?:\s*\[([^\]]*)\])?"  # (6) existing options, without brackets
        r"\s*;"  # trailing semicolon
    )

    for i, line in enumerate(lines):
        # Insert nanopb import right after the last existing import line.
        # Only do this when there IS an existing import (last_import_idx >= 0);
        # if there are no imports we fall through to the syntax-line fallback below.
        if not import_added and last_import_idx >= 0 and i == last_import_idx + 1:
            result.append(f'import "{nanopb_import_path}";')
            import_added = True

        # --- Track message/oneof/enum nesting ---
        m = message_re.match(line)
        if m:
            context_stack.append(("message", m.group(2)))
            result.append(line)
            continue

        m = oneof_re.match(line)
        if m:
            context_stack.append(("oneof", m.group(2)))
            result.append(line)
            continue

        m = enum_re.match(line)
        if m:
            context_stack.append(("enum", m.group(2)))
            result.append(line)
            continue

        if close_re.match(line) and context_stack:
            context_stack.pop()
            result.append(line)
            continue

        # Skip field injection inside enum bodies (enum values look like fields
        # but should not have nanopb options added)
        in_enum = bool(context_stack) and context_stack[-1][0] == "enum"

        # --- Try to match and modify a field declaration ---
        m = field_re.match(line)
        if m and not in_enum:
            indent = m.group(1)
            qualifier = m.group(2) or ""
            ftype = m.group(3)
            fname = m.group(4)
            fnum = m.group(5)
            existing_opts = m.group(6) or ""

            # Collect applicable nanopb options (wildcard < specific)
            extra: Dict[str, Any] = {}

            # 1. Wildcard: any field with this name in this proto file
            if fname in wildcard:
                extra.update(wildcard[fname])

            # 2. Specific: check all keys whose last element is fname and whose
            #    preceding path matches the current message context
            for key, opts in specific.items():
                if key[-1] == fname:
                    msg_path = key[:-1]
                    if message_path_matches(context_stack, msg_path):
                        extra.update(opts)
                        break

            if extra:
                nanopb_str = format_nanopb_opts(extra)
                if existing_opts.strip():
                    opts_block = f"[{existing_opts}, {nanopb_str}]"
                else:
                    opts_block = f"[{nanopb_str}]"
                qual = qualifier.rstrip()
                sep = " " if qual else ""
                line = f"{indent}{qual}{sep}{ftype} {fname} = {fnum} {opts_block};"

        result.append(line)

    # Edge case: if there were no import lines, add nanopb import after syntax line
    if not import_added:
        for i, line in enumerate(result):
            if line.strip().startswith("syntax") and line.strip().endswith(";"):
                result.insert(i + 1, f'import "{nanopb_import_path}";')
                break

    return "\n".join(result)


def main() -> int:
    """Parse an .options file and inject its constraints into a .proto file in-place."""
    if len(sys.argv) != 3:
        print(
            f"Usage: {sys.argv[0]} <options_file> <proto_file>",
            file=sys.stderr,
        )
        return 1

    opts_path = Path(sys.argv[1])
    proto_path = Path(sys.argv[2])

    if not opts_path.exists():
        print(f"Options file not found: {opts_path}", file=sys.stderr)
        return 1

    if not proto_path.exists():
        print(f"Proto file not found: {proto_path}", file=sys.stderr)
        return 1

    specific, wildcard = parse_options_file(opts_path)
    total = len(specific) + len(wildcard)

    if total == 0:
        print(f"  [{opts_path.name}] No injectable options found, skipping.")
        return 0

    content = proto_path.read_text(encoding="utf-8")

    # After regen-protobufs.sh's sed fixup, the nanopb import path is:
    nanopb_import_path = "meshtastic/protobuf/nanopb.proto"

    modified = inject_into_proto(content, specific, wildcard, nanopb_import_path)
    proto_path.write_text(modified, encoding="utf-8")

    print(
        f"  [{opts_path.name}] Injected {len(specific)} specific + "
        f"{len(wildcard)} wildcard option(s) into {proto_path.name}"
    )
    return 0


if __name__ == "__main__":
    sys.exit(main())