extra_script.py

#!/usr/bin/env python3
"""
PlatformIO pre-build script to generate protobuf files and dynamic modules for plugins.
"""

import os
import re
import subprocess
import sys
from pathlib import Path

Import("env")  # type: ignore[name-defined]  # noqa: F821

project_dir = env["PROJECT_DIR"]  # type: ignore[name-defined]  # noqa: F821
pio_env = env.get("PIOENV")  # type: ignore[name-defined]  # noqa: F821

if not pio_env:
    print("Meshtastic-Plugins: PIOENV not found, skipping plugin generation")
    sys.exit(0)

libdeps_dir = os.path.join(project_dir, ".pio", "libdeps", pio_env)

if not os.path.exists(libdeps_dir):
    print(f"Meshtastic-Plugins: Library dependencies directory not found: {libdeps_dir}")
    print("Meshtastic-Plugins: Plugins will be scanned after PlatformIO installs dependencies")
    # Generate empty DynamicModules.cpp
    generated_dir = os.path.join(project_dir, "src", "mesh", "generated")
    os.makedirs(generated_dir, exist_ok=True)
    cpp_file = os.path.join(generated_dir, "DynamicModules.cpp")
    with open(cpp_file, "w", encoding="utf-8") as f:
        f.write("// src/mesh/generated/DynamicModules.cpp (GENERATED - DO NOT EDIT)\n")
        f.write("// This file is auto-generated by meshtastic-plugins. Do not edit manually.\n\n")
        f.write('#include "DebugConfiguration.h"\n\n')
        f.write("void init_dynamic_modules() {\n")
        f.write("    // No plugins installed\n")
        f.write("}\n")
    sys.exit(0)


def ensure_nanopb():
    """Ensure nanopb package is installed."""
    try:
        import nanopb  # noqa: F401
        # Check if nanopb_generator command exists
        result = subprocess.run(
            ["nanopb_generator", "--version"],
            capture_output=True,
            text=True,
            timeout=5,
        )
        if result.returncode == 0:
            return True
    except (ImportError, FileNotFoundError, subprocess.TimeoutExpired):
        pass

    # Install nanopb
    print("Installing nanopb package...")
    try:
        env.Execute("$PYTHONEXE -m pip install 'nanopb>=0.4.9'")  # type: ignore[name-defined]  # noqa: F821
        return True
    except Exception as e:
        print(f"Warning: Failed to install nanopb: {e}", file=sys.stderr)
        return False


def scan_plugins():
    """
    Scan for plugins in PlatformIO libdeps directory.

    Returns:
        List of tuples (plugin_name, plugin_path, include_path, src_path, proto_files)
    """
    plugins = []

    if not os.path.exists(libdeps_dir) or not os.path.isdir(libdeps_dir):
        return plugins

    for lib_dir_name in os.listdir(libdeps_dir):
        lib_dir = os.path.join(libdeps_dir, lib_dir_name)
        if not os.path.isdir(lib_dir):
            continue

        # Check for include/plugin.h (plugin marker)
        plugin_h = os.path.join(lib_dir, "include", "plugin.h")
        if not os.path.exists(plugin_h):
            continue

        include_path = os.path.join(lib_dir, "include")
        src_path = os.path.join(lib_dir, "src")
        
        # Scan for .proto files recursively in src directory
        proto_files = []
        if os.path.isdir(src_path):
            for root, dirs, files in os.walk(src_path):
                # Skip hidden directories
                dirs[:] = [d for d in dirs if not d.startswith(".")]
                for file in files:
                    if file.endswith(".proto"):
                        proto_files.append(os.path.join(root, file))

        plugins.append((lib_dir_name, lib_dir, include_path, src_path, proto_files))

    return plugins


def generate_protobuf_files(proto_file, options_file=None, output_dir=None):
    """
    Generate protobuf C++ files using nanopb.

    Args:
        proto_file: Path to the .proto file
        options_file: Optional path to .options file
        output_dir: Optional output directory (defaults to proto file directory)

    Returns:
        bool: True if successful, False otherwise
    """
    proto_file = os.path.abspath(proto_file)
    if not os.path.exists(proto_file):
        print(f"Error: Proto file not found: {proto_file}")
        return False

    proto_dir = os.path.dirname(proto_file)
    proto_basename = os.path.basename(proto_file)
    proto_name = os.path.splitext(proto_basename)[0]

    if output_dir is None:
        output_dir = proto_dir
    else:
        output_dir = os.path.abspath(output_dir)

    os.makedirs(output_dir, exist_ok=True)

    # Auto-detect options file if not specified
    if options_file is None:
        candidate_options = os.path.join(proto_dir, f"{proto_name}.options")
        if os.path.exists(candidate_options):
            options_file = candidate_options

    print(f"Generating protobuf files from {proto_basename}...")

    cmd = [
        "nanopb_generator",
        "-D",
        output_dir,
        "-I",
        proto_dir,
        "-S",
        ".cpp",
        proto_file,
    ]

    try:
        result = subprocess.run(cmd, cwd=proto_dir, check=True, capture_output=True, text=True)
        return True
    except subprocess.CalledProcessError as e:
        print(f"Error generating protobufs: {e}")
        if e.stdout:
            print(e.stdout)
        if e.stderr:
            print(e.stderr)
        return False
    except FileNotFoundError:
        print("Error: nanopb_generator command not found. Make sure nanopb is installed.")
        return False


def generate_all_protobuf_files(plugins, verbose=True):
    """
    Generate protobuf files for all proto files found in plugins.

    Args:
        plugins: List of plugin tuples from scan_plugins()
        verbose: Whether to print status messages

    Returns:
        Tuple of (success_count, total_count)
    """
    success_count = 0
    total_count = 0

    for plugin_name, plugin_path, include_path, src_path, proto_files in plugins:
        for proto_file in proto_files:
            total_count += 1
            proto_basename = os.path.basename(proto_file)
            proto_dir = os.path.dirname(proto_file)
            proto_name = os.path.splitext(proto_basename)[0]

            options_file = os.path.join(proto_dir, f"{proto_name}.options")
            options_path = options_file if os.path.exists(options_file) else None

            if verbose:
                print(f"Meshtastic-Plugins: Processing {proto_basename} from {plugin_name}...")

            if generate_protobuf_files(proto_file, options_path, proto_dir):
                success_count += 1
            elif verbose:
                print(f"Meshtastic-Plugins: Failed to generate protobuf files for {proto_basename}")

    return success_count, total_count


def generate_dynamic_modules(project_dir, plugins, verbose=True):
    """
    Generate DynamicModules.cpp from plugin headers that contain #pragma MPM_MODULE.

    Args:
        project_dir: Root directory of the project
        plugins: List of plugin tuples from scan_plugins()
        verbose: Whether to print status messages

    Returns:
        bool: True if successful, False otherwise
    """
    generated_dir = os.path.join(project_dir, "src", "mesh", "generated")
    cpp_file = os.path.join(generated_dir, "DynamicModules.cpp")

    os.makedirs(generated_dir, exist_ok=True)

    # Collect module registrations from plugin headers
    module_registrations = []
    pragma_pattern = re.compile(
        r'#pragma\s+MPM_MODULE\s*\(\s*(\w+)(?:\s*,\s*(\w+))?\s*\)'
    )

    for plugin_name, plugin_path, include_path, src_path, proto_files in plugins:
        # Scan include/plugin.h for #pragma MPM_MODULE
        plugin_h = os.path.join(include_path, "plugin.h")
        if not os.path.exists(plugin_h):
            continue

        try:
            with open(plugin_h, "r", encoding="utf-8") as f:
                content = f.read()
                for match in pragma_pattern.finditer(content):
                    class_name = match.group(1)
                    variable_name = match.group(2) if match.group(2) else None
                    # Use library directory name directly in include path
                    header_filename = f"{plugin_name}/include/plugin.h"
                    module_registrations.append({
                        "plugin_name": plugin_name,
                        "class_name": class_name,
                        "header_filename": header_filename,
                        "variable_name": variable_name,
                    })
                    if verbose:
                        var_info = f" -> {variable_name}" if variable_name else ""
                        print(f"Meshtastic-Plugins: Found module {class_name} in {plugin_name}{var_info}")
        except Exception as e:
            if verbose:
                print(f"Meshtastic-Plugins: Warning: Failed to read {plugin_h}: {e}")

    # Generate the .cpp file
    try:
        with open(cpp_file, "w", encoding="utf-8") as f:
            f.write("// src/mesh/generated/DynamicModules.cpp (GENERATED - DO NOT EDIT)\n")
            f.write("// This file is auto-generated by meshtastic-plugins. Do not edit manually.\n\n")

            f.write('#include "DebugConfiguration.h"\n')

            # Collect unique header includes
            header_includes = set()
            for reg in module_registrations:
                header_includes.add(reg["header_filename"])

            # Write includes
            if header_includes:
                f.write("\n")
            for header in sorted(header_includes):
                f.write(f'#include "{header}"\n')

            if header_includes:
                f.write("\n")

            # Write function implementation
            f.write("void init_dynamic_modules() {\n")

            if module_registrations:
                current_plugin = None
                for reg in module_registrations:
                    if reg["plugin_name"] != current_plugin:
                        if current_plugin is not None:
                            f.write("\n")
                        f.write(f'    // Plugin: {reg["plugin_name"]}\n')
                        current_plugin = reg["plugin_name"]

                    f.write(f'    LOG_DEBUG("Initializing module {reg["class_name"]} from plugin {reg["plugin_name"]}");\n')

                    if reg["variable_name"]:
                        f.write(f'    {reg["variable_name"]} = new {reg["class_name"]}();\n')
                    else:
                        f.write(f'    new {reg["class_name"]}();\n')
            else:
                f.write("    // No plugins installed\n")

            f.write("}\n")

        if verbose:
            print(f"Meshtastic-Plugins: Generated DynamicModules.cpp with {len(module_registrations)} module(s)")
        return True

    except Exception as e:
        print(f"Meshtastic-Plugins: Error generating DynamicModules.cpp: {e}", file=sys.stderr)
        return False


# Main execution
plugins = scan_plugins()

if not plugins:
    print("Meshtastic-Plugins: No plugins found")
    # Still generate empty DynamicModules.cpp
    generate_dynamic_modules(project_dir, [], verbose=True)
    sys.exit(0)

print(f"Meshtastic-Plugins: Found {len(plugins)} plugin(s)")

# Ensure nanopb is available
if not ensure_nanopb():
    print("Warning: nanopb not available, skipping protobuf generation", file=sys.stderr)
else:
    # Generate protobuf files
    print("Meshtastic-Plugins: Generating protobuf files...")
    success_count, total_count = generate_all_protobuf_files(plugins, verbose=True)
    print(f"Meshtastic-Plugins: Completed - {success_count}/{total_count} protobuf file(s) generated successfully")

# Generate dynamic modules
print("Meshtastic-Plugins: Generating dynamic modules...")
if not generate_dynamic_modules(project_dir, plugins, verbose=True):
    print("Warning: Failed to generate dynamic modules", file=sys.stderr)