# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: OpenMDW-1.1

"""Shared helpers for the standalone Action policy serving scripts.

Both ``cosmos_framework.scripts.action_policy_server_libero`` (HTTP) and
``cosmos_framework.scripts.action_policy_server_robolab`` (WebSocket) reuse the
helpers defined here:

- single-rank distributed init for FSDP-aware model loading;
- local IP / free-port discovery for log messages and PG bring-up;
- the default output directory shared across both servers;
- the ``OmniSetupArgs`` subclass that disables runtime EMA for frozen configs.

Keeping these in one module avoids ``action_policy_server_robolab`` importing
from ``action_policy_server_libero`` purely to share runtime utilities.
"""

import os
import socket
from pathlib import Path

import torch
from torch import distributed as dist

from cosmos_framework.inference.args import OmniSetupArgs
from cosmos_framework.inference.common.args import ConfigFileType
from cosmos_framework.utils import log

DEFAULT_FALLBACK_OUTPUT_DIR = Path("/tmp/cosmos3_action_server")


def _get_free_local_port() -> int:
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        s.bind(("127.0.0.1", 0))
        return int(s.getsockname()[1])


def get_local_ip() -> str:
    """Get the local IP address of this machine."""
    try:
        # Connect to an external address to determine the local IP
        with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as s:
            s.connect(("8.8.8.8", 80))
            return str(s.getsockname()[0])
    except Exception:
        return socket.gethostbyname(socket.gethostname())


def maybe_init_distributed() -> None:
    """
    Initialize a one-rank process group when launched outside torchrun so FSDP
    and distributed utilities have the process-group state they expect.

    ``init_script()`` already inits the PG when ``WORLD_SIZE>1`` (under
    torchrun); this function fills in the single-process case used by the
    standalone server.
    """
    if not dist.is_available() or dist.is_initialized():
        return

    world_size_env = os.getenv("WORLD_SIZE")
    rank_env = os.getenv("RANK")
    local_rank = int(os.getenv("LOCAL_RANK", "0"))

    backend = "nccl" if torch.cuda.is_available() else "gloo"

    if world_size_env is not None and rank_env is not None:
        if torch.cuda.is_available():
            torch.cuda.set_device(local_rank)
        dist.init_process_group(backend=backend, init_method="env://")
        return

    if torch.cuda.is_available():
        torch.cuda.set_device(0)
    port = _get_free_local_port()
    dist.init_process_group(backend=backend, init_method=f"tcp://127.0.0.1:{port}", rank=0, world_size=1)


class _ActionPolicyServerSetupArgs(OmniSetupArgs):
    """Server-local setup args that avoid instantiating runtime EMA for frozen configs."""

    def load_model_config_dict(self) -> dict:
        model_dict = super().load_model_config_dict()
        model_dict.setdefault("config", {}).setdefault("ema", {})["enabled"] = False
        return model_dict


def disable_runtime_ema_for_frozen_config(setup_args: OmniSetupArgs) -> OmniSetupArgs:
    """Use server-local setup args to instantiate frozen configs without runtime EMA."""
    if setup_args.config_file_type == ConfigFileType.MODULE:
        return setup_args

    log.info("[action-server] disabled runtime EMA for frozen config model load")
    return _ActionPolicyServerSetupArgs.model_validate(setup_args.model_dump())