ReMe/reme/utils/service_utils.py

"""Service discovery utilities."""

import asyncio
import socket
import sys

import psutil

from ..constants import (
    REME_DEFAULT_HOST,
    REME_DEFAULT_PORT,
    normalize_connect_host,
)


async def find_reme(host: str, port: int) -> str:
    """Probe host:port. Returns 'reme', 'occupied', or 'free'."""
    from ..components.client.http_client import HttpClient

    try:
        async with HttpClient(host=normalize_connect_host(host), port=port, timeout=2.0) as client:
            async for _ in client(action="health_check"):
                break
        return "reme"
    except Exception:
        pass
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        try:
            s.bind((host, port))
            return "free"
        except OSError:
            return "occupied"


def _pid_on_port(port: int) -> int | None:
    """PID listening on TCP ``port``, or None. Cross-platform via psutil.

    Iterates per-process rather than calling the system-wide
    ``psutil.net_connections()`` — the latter needs root on macOS, while
    per-process connection enumeration works without elevation for
    processes the current user owns (which reme's own server always is).
    """
    for proc in psutil.process_iter(["pid"]):
        try:
            for conn in proc.net_connections(kind="tcp"):
                if conn.status == psutil.CONN_LISTEN and conn.laddr and conn.laddr.port == port:
                    return proc.info["pid"]
        except (psutil.AccessDenied, psutil.NoSuchProcess, psutil.ZombieProcess):
            continue
    return None


def _start_arguments(cmdline: list[str]) -> list[str] | None:
    """Identify a ReMe CLI process, excluding supervisors containing child argv."""
    if not cmdline:
        return None

    def command_name(value: str) -> str:
        return value.replace("\\", "/").rsplit("/", 1)[-1].lower()

    executable = command_name(cmdline[0])
    if executable in {"reme", "reme.exe"}:
        action_index = 1
    elif executable.startswith(("python", "pypy")):
        action_index = next(
            (
                index
                for index in range(2, len(cmdline))
                if command_name(cmdline[index - 1]) in {"reme", "reme.exe", "reme.py", "reme.reme"}
            ),
            len(cmdline),
        )
        if "-c" in cmdline[:action_index]:
            return None
    else:
        return None
    if action_index >= len(cmdline) or cmdline[action_index] not in {"start", "-start", "--start"}:
        return None
    return [argument for argument in cmdline[action_index + 1 :] if "=" in argument]


def _scan_reme_procs() -> list[tuple[int, str, int]]:
    """List running 'reme ... start' processes as (pid, host, port)."""
    from ..config import parse_kwargs, resolve_app_config

    procs: list[tuple[int, str, int]] = []
    for proc in psutil.process_iter(["pid", "cmdline"]):
        try:
            cmdline = proc.info["cmdline"] or []
        except (psutil.NoSuchProcess, psutil.AccessDenied):
            continue
        arguments = _start_arguments(cmdline)
        if arguments is None:
            continue
        try:
            overrides = parse_kwargs(*arguments)
        except ValueError:
            continue
        try:
            config = resolve_app_config(log_config=False, **overrides)
        except (OSError, ValueError):
            # A different process's relative config file may be unavailable.
            config = overrides
        service = config.get("service") or {}
        host = service.get("host") or REME_DEFAULT_HOST
        port_value = str(service.get("port", REME_DEFAULT_PORT))
        port = int(port_value) if port_value.isdigit() else REME_DEFAULT_PORT
        procs.append((proc.info["pid"], host, port))
    return procs


def _reme_start_argv() -> list[list[str]]:
    """Return the `key=value` start-args of each running `reme ... start` process.

    These cmdline tokens are the authoritative record of how the server was
    *actually* launched — including ``service.*`` overrides that never touched
    any config file — so replaying them reproduces the server's own config.
    """
    argvs: list[list[str]] = []
    for proc in psutil.process_iter(["cmdline"]):
        try:
            cmdline = proc.info["cmdline"] or []
        except (psutil.NoSuchProcess, psutil.AccessDenied):
            continue
        arguments = _start_arguments(cmdline)
        if arguments is not None:
            argvs.append(arguments)
    return argvs


def running_app_config() -> dict | None:
    """Resolve the full config of a running ReMe by replaying its start args.

    Reads the live ``reme start ...`` process cmdline and re-runs the same
    ``resolve_app_config`` the server used. Returning the full config preserves
    enabled plugins as well as service connection settings for CLI clients.
    Returns ``None`` when no running reme is found or its args can't be parsed.
    """
    from ..config import parse_args, resolve_app_config

    for argv in _reme_start_argv():
        try:
            _, kwargs = parse_args("start", *argv)
            config = resolve_app_config(log_config=False, **kwargs)
        except ValueError:
            continue
        if isinstance(config, dict):
            return config
    return None


def running_service_config() -> dict | None:
    """Return only the service section of the running ReMe configuration."""
    config = running_app_config()
    if config is None:
        return None
    service = config.get("service")
    return service if isinstance(service, dict) else None


async def locate_reme() -> tuple[str, int, int | None] | None:
    """Find a running reme: try default port, then scanned processes."""
    if await find_reme(REME_DEFAULT_HOST, REME_DEFAULT_PORT) == "reme":
        return REME_DEFAULT_HOST, REME_DEFAULT_PORT, _pid_on_port(REME_DEFAULT_PORT)
    for pid, host, port in _scan_reme_procs():
        if await find_reme(host, port) == "reme":
            return normalize_connect_host(host), port, pid
    return None


def precheck_start(svc_config: dict | None) -> bool:
    """Pre-flight check for `start`: False if reme is up, exits 1 on port conflict."""
    host = (svc_config or {}).get("host") or REME_DEFAULT_HOST
    port = (svc_config or {}).get("port") or REME_DEFAULT_PORT
    port = int(port)
    status = asyncio.run(find_reme(host, port))
    if status == "reme":
        print(f"reme already running at {normalize_connect_host(host)}:{port}")
        return False
    if status == "occupied":
        print(
            f"port {port} occupied. Start on another port: reme start service.port=<other_port>",
            file=sys.stderr,
        )
        sys.exit(1)
    return True


def cli_find_reme() -> None:
    """Handle `reme find_reme`: print HOST/PORT/PID or a hint to start reme."""
    found = asyncio.run(locate_reme())
    if not found:
        print("reme not started. Try: reme start", file=sys.stderr)
        sys.exit(1)
    host, port, pid = found
    print(f"HOST={host} PORT={port} PID={pid or 'unknown'}")