#!/usr/bin/env python3
"""Generate the Python Shard-protocol stubs from `shard_runtime.proto`.

The `.proto` file is the contract; the Python modules under
`meshnet_node/native_protocol/generated/` are build output that happens to be
committed. They are committed so that installing the node does not require a
protoc toolchain, and `--check` exists so a committed stub can never silently
drift from the schema it claims to implement.

Usage::

    python scripts/generate_native_protocol.py           # regenerate in place
    python scripts/generate_native_protocol.py --check   # fail if out of date

C++ stubs are *not* generated here. They are build artifacts produced by CMake
(`packages/node/native/CMakeLists.txt`) into the build tree, because a C++ build
already requires a toolchain and nothing is gained by committing them.
"""

from __future__ import annotations

import argparse
from importlib import metadata
import pathlib
import shutil
import subprocess
import sys
import tempfile

REPO_ROOT = pathlib.Path(__file__).resolve().parent.parent
PROTO_DIR = REPO_ROOT / "packages/node/native/proto"
PROTO_FILE = PROTO_DIR / "shard_runtime.proto"
OUT_DIR = REPO_ROOT / "packages/node/meshnet_node/native_protocol/generated"

# Regenerating with a different protoc emits different gencode headers, so the
# generator version is part of the contract and `--check` would catch a drift.
REQUIRED_GRPCIO_TOOLS = "1.82.1"

_HEADER = "# Generated by scripts/generate_native_protocol.py. Do not edit.\n"


def _require_grpcio_tools_version() -> None:
    try:
        actual = metadata.version("grpcio-tools")
    except metadata.PackageNotFoundError:
        sys.exit(
            "grpc_tools is required to generate stubs:\n"
            f"    pip install grpcio-tools=={REQUIRED_GRPCIO_TOOLS}"
        )
    if actual != REQUIRED_GRPCIO_TOOLS:
        sys.exit(
            "wrong grpcio-tools version for deterministic generation: "
            f"found {actual}, require {REQUIRED_GRPCIO_TOOLS}\n"
            f"    pip install --upgrade grpcio-tools=={REQUIRED_GRPCIO_TOOLS}"
        )


def _generate(into: pathlib.Path) -> None:
    """Run the exactly pinned protoc, writing generated modules into `into`."""
    _require_grpcio_tools_version()
    try:
        from grpc_tools import protoc
    except ImportError:  # pragma: no cover - inconsistent/broken installation
        sys.exit(
            "grpcio-tools metadata exists but grpc_tools cannot be imported; reinstall it:\n"
            f"    pip install --force-reinstall grpcio-tools=={REQUIRED_GRPCIO_TOOLS}"
        )

    into.mkdir(parents=True, exist_ok=True)
    # grpc_tools bundles protoc and the well-known types, so generation needs no
    # system protoc and produces identical output on any machine.
    well_known = pathlib.Path(protoc.__file__).parent / "_proto"
    args = [
        "protoc",
        f"--proto_path={PROTO_DIR}",
        f"--proto_path={well_known}",
        f"--python_out={into}",
        f"--pyi_out={into}",
        f"--grpc_python_out={into}",
        str(PROTO_FILE),
    ]
    if protoc.main(args) != 0:
        sys.exit("protoc failed")

    # protoc emits `import shard_runtime_pb2` — a bare top-level import that only
    # resolves if the generated directory happens to be on sys.path. Rewrite it
    # to a relative import so the package works as an installed package.
    grpc_module = into / "shard_runtime_pb2_grpc.py"
    text = grpc_module.read_text()
    text = text.replace(
        "import shard_runtime_pb2 as shard__runtime__pb2",
        "from . import shard_runtime_pb2 as shard__runtime__pb2",
    )
    grpc_module.write_text(text)

    (into / "__init__.py").write_text(
        _HEADER + '"""Generated protobuf/gRPC stubs for the native Shard protocol."""\n'
    )


def _tracked_files(directory: pathlib.Path) -> dict[str, bytes]:
    return {
        path.name: path.read_bytes()
        for path in sorted(directory.iterdir())
        if path.is_file() and path.suffix in {".py", ".pyi"}
    }


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--check",
        action="store_true",
        help="verify the committed stubs match the .proto instead of rewriting them",
    )
    args = parser.parse_args()

    if args.check:
        with tempfile.TemporaryDirectory() as tmp:
            fresh = pathlib.Path(tmp) / "generated"
            _generate(fresh)
            if not OUT_DIR.is_dir():
                print(f"generated stubs are missing: {OUT_DIR}", file=sys.stderr)
                return 1
            if _tracked_files(fresh) != _tracked_files(OUT_DIR):
                print(
                    "generated stubs are out of date with shard_runtime.proto.\n"
                    "Run: python scripts/generate_native_protocol.py",
                    file=sys.stderr,
                )
                return 1
        print("generated stubs are up to date")
        return 0

    if OUT_DIR.exists():
        shutil.rmtree(OUT_DIR)
    _generate(OUT_DIR)
    print(f"wrote {OUT_DIR.relative_to(REPO_ROOT)}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
