#!/usr/bin/env python3
"""Generate RPC method registry from upstream proto files."""

from __future__ import annotations

import re
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
PROTO_SOURCES = [
    ("omni.resources.ResourceService", ROOT / "refs/libs/omni/client/api/omni/resources/resources.proto"),
    ("management.ManagementService", ROOT / "refs/libs/omni/client/api/omni/management/management.proto"),
    ("oidc.OIDCService", ROOT / "refs/libs/omni/client/api/omni/oidc/oidc.proto"),
    ("auth.AuthService", ROOT / "refs/libs/omni/frontend/src/api/omni/auth/auth.proto"),
]
OUT = ROOT / "src/omni/generated/methods.py"

RPC_RE = re.compile(r"^\s*rpc\s+(\w+)\([^)]*\)\s+returns\s+\((stream\s+)?", re.MULTILINE)


def const_name(service: str, method: str) -> str:
    service_part = service.split(".")[-1].replace("Service", "").upper()
    tokens = re.sub(r"([a-z0-9])([A-Z])", r"\1_\2", method).upper()
    return f"{service_part}_{tokens}"


def generate() -> str:
    lines: list[str] = []
    lines.append('"""Generated RPC method registry.\n+\n+This file is generated by scripts/generate_methods.py.\n+"""')
    lines.append("")
    lines.append("from __future__ import annotations")
    lines.append("")
    lines.append("from omni.http.methods import RpcMethod")
    lines.append("")

    all_names: list[str] = []
    for service, path in PROTO_SOURCES:
        content = path.read_text(encoding="utf-8")
        for method, stream_mark in RPC_RE.findall(content):
            cname = const_name(service, method)
            stream = bool(stream_mark)
            idempotent = method.startswith(("Get", "List", "Read", "Validate", "Kubeconfig", "Talosconfig", "Omniconfig", "Controllers", "DependencyGraph", "Watch", "Await"))
            lines.append(
                f'{cname} = RpcMethod("{service}", "{method}", stream={str(stream)}, idempotent={str(idempotent)})'
            )
            all_names.append(cname)
        lines.append("")

    lines.append("ALL_METHODS = {")
    lines.append("    method.key: method")
    lines.append("    for method in (")
    for name in all_names:
        lines.append(f"        {name},")
    lines.append("    )")
    lines.append("}")

    return "\n".join(lines) + "\n"


def main() -> None:
    OUT.parent.mkdir(parents=True, exist_ok=True)
    OUT.write_text(generate(), encoding="utf-8")
    print(f"wrote {OUT}")


if __name__ == "__main__":
    main()
