# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import asyncio
import json
import os
import shlex
from functools import wraps
from glob import glob
from os import makedirs
from os.path import exists
from pathlib import Path
from shutil import rmtree
from signal import SIGINT
from subprocess import Popen, TimeoutExpired
from threading import Thread
from time import sleep, time
from typing import Dict, List, Optional, Tuple

import rich
import uvicorn
from devtools import pprint
from omegaconf import DictConfig, OmegaConf
from pydantic import Field
from rich.markup import escape
from rich.table import Table
from tqdm.auto import tqdm

from nemo_gym import PARENT_DIR, ROOT_DIR
from nemo_gym.cli.setup_command import run_command, setup_env_command
from nemo_gym.config_types import BaseNeMoGymCLIConfig, ConfigError
from nemo_gym.global_config import (
    DRY_RUN_KEY_NAME,
    JSON_OUTPUT_KEY_NAME,
    NEMO_GYM_CONFIG_DICT_ENV_VAR_NAME,
    NEMO_GYM_CONFIG_PATH_ENV_VAR_NAME,
    NEMO_GYM_RESERVED_TOP_LEVEL_KEYS,
    GlobalConfigDictParser,
    GlobalConfigDictParserConfig,
    get_global_config_dict,
)
from nemo_gym.registry import discover_environments
from nemo_gym.server_status import StatusCommand
from nemo_gym.server_utils import (
    HEAD_SERVER_KEY_NAME,
    HeadServer,
    ServerClient,
    ServerInstanceDisplayConfig,
    ServerStatus,
    initialize_ray,
)


# Grace period after SIGINT before escalating to SIGKILL. Kept short so Ctrl-C is responsive.
_GRACEFUL_SHUTDOWN_TIMEOUT_SEC: int = 1
# Grace period after SIGKILL for the kernel to reap the child and avoid <defunct> entries.
_FORCE_KILL_REAP_TIMEOUT_SEC: int = 2


def exit_cleanly_on_config_error(fn):
    """Decorator: turn user-facing ConfigError into a clean message + non-zero exit.

    Config mistakes (missing/typo'd config_paths, malformed config_paths, nothing configured to
    run) should fail fast with an actionable message, not a Python traceback. Unexpected errors
    still propagate normally.
    """

    @wraps(fn)
    def wrapper(*args, **kwargs):
        try:
            return fn(*args, **kwargs)
        except ConfigError as e:
            # escape() so '[...]' in the message (e.g. config_paths examples) isn't eaten as rich markup.
            rich.print(f"[red]Error:[/red] {escape(str(e))}")
            raise SystemExit(1)

    return wrapper


def _resolve_server_dir(rel_path: Path) -> Path:
    """Resolve a relative server dir (e.g. ``resources_servers/<name>``) to an absolute path.

    Checks the current working directory first (a user's local server), then falls back to the Gym
    install root (``PARENT_DIR``) where built-in servers live in both editable and wheel installs.
    This lets ``gym env test`` find and run built-in servers from any cwd, not just a repo checkout.
    """
    cwd_path = Path.cwd() / rel_path
    if (cwd_path / "requirements.txt").exists() or (cwd_path / "pyproject.toml").exists():
        return cwd_path
    return PARENT_DIR / rel_path


class RunConfig(BaseNeMoGymCLIConfig):
    """
    Start NeMo Gym servers for agents, models, and resources.

    Examples:

    ```bash
    config_paths="resources_servers/example_single_tool_call/configs/example_single_tool_call.yaml,\\
    responses_api_models/openai_model/configs/openai_model.yaml"
    gym env start "+config_paths=[${config_paths}]"
    ```
    """

    entrypoint: str = Field(
        description="Entrypoint for this command. This must be a relative path with 2 parts. Should look something like `responses_api_agents/simple_agent`."
    )


class TestConfig(RunConfig):
    """
    Test a specific server module by running its pytest suite and optionally validating example data.

    Examples:

    ```bash
    gym env test +entrypoint=resources_servers/example_single_tool_call
    ```
    """

    should_validate_data: bool = Field(
        default=False,
        description="Whether or not to validate the example data (examples, metrics, rollouts, etc) for this server.",
    )

    _dir_path: Path  # initialized in model_post_init

    def model_post_init(self, context):  # pragma: no cover
        self._dir_path = Path(self.entrypoint)
        assert not self.dir_path.is_absolute()

        return super().model_post_init(context)

    @property
    def dir_path(self) -> Path:
        return self._dir_path

    @property
    def resolved_dir_path(self) -> Path:
        """Absolute server dir resolved against the cwd, then the Gym install root.

        Use this for filesystem access (reading data, running the suite); use ``dir_path`` (the
        relative entrypoint) for display and example commands shown to the user.
        """
        return _resolve_server_dir(self._dir_path)


class RunHelper:  # pragma: no cover
    _head_server: uvicorn.Server
    _head_server_thread: Thread
    _head_server_instance: HeadServer

    _processes: Dict[str, Popen]
    _server_instance_display_configs: List[ServerInstanceDisplayConfig]
    _server_client: ServerClient

    def start(self, global_config_dict_parser_config: GlobalConfigDictParserConfig) -> None:
        global_config_dict = get_global_config_dict(global_config_dict_parser_config=global_config_dict_parser_config)

        # Fail fast before starting Ray if nothing is configured to run (covers env run and the
        # e2e rollout-collection path, which both start servers via this method).
        GlobalConfigDictParser().raise_on_no_server_instances(global_config_dict)

        # Initialize Ray cluster in the main process
        # Note: This function will modify the global config dict - update `ray_head_node_address`
        initialize_ray()

        # Assume Nemo Gym Run is for a single agent.
        escaped_config_dict_yaml_str = shlex.quote(OmegaConf.to_yaml(global_config_dict))

        # We always run the head server in this `run` command.
        self._head_server, self._head_server_thread, self._head_server_instance = HeadServer.run_webserver()

        top_level_paths = [k for k in global_config_dict.keys() if k not in NEMO_GYM_RESERVED_TOP_LEVEL_KEYS]

        self._processes: Dict[str, Popen] = dict()
        self._server_instance_display_configs: List[ServerInstanceDisplayConfig] = []

        start_time = time()

        # TODO there is a better way to resolve this that uses nemo_gym/global_config.py::ServerInstanceConfig
        for top_level_path in top_level_paths:
            server_config_dict = global_config_dict[top_level_path]
            if not isinstance(server_config_dict, DictConfig):
                continue

            first_key = list(server_config_dict)[0]
            server_config_dict = server_config_dict[first_key]
            if not isinstance(server_config_dict, DictConfig):
                continue
            second_key = list(server_config_dict)[0]
            server_config_dict = server_config_dict[second_key]
            if not isinstance(server_config_dict, DictConfig):
                continue

            if "entrypoint" not in server_config_dict:
                continue

            # TODO: This currently only handles relative entrypoints. Later on we can resolve the absolute path.
            entrypoint_fpath = Path(server_config_dict.entrypoint)
            assert not entrypoint_fpath.is_absolute()

            # Resolve cwd-first (a local server), else the install location for built-ins.
            dir_path = _resolve_server_dir(Path(first_key, second_key))

            command = f"""{setup_env_command(dir_path, global_config_dict, top_level_path)} \\
    && {NEMO_GYM_CONFIG_DICT_ENV_VAR_NAME}={escaped_config_dict_yaml_str} \\
    {NEMO_GYM_CONFIG_PATH_ENV_VAR_NAME}={shlex.quote(top_level_path)} \\
    python {str(entrypoint_fpath)}"""

            process = run_command(command, dir_path, server_name=top_level_path)
            self._processes[top_level_path] = process
            # In dry run mode, wait for each setup command to finish before starting the next.
            # This installs uv virtual environments serially, which significantly reduces uv
            # cache size. For Nemotron's set of environments, parallel installation can produce
            # a cache 10-20GB larger than serial installation.
            if global_config_dict[DRY_RUN_KEY_NAME]:
                print("DRY_RUN enabled: setup commands are run serially")
                process.communicate()

            host = server_config_dict.get("host")
            port = server_config_dict.get("port")

            self._server_instance_display_configs.append(
                ServerInstanceDisplayConfig(
                    process_name=top_level_path,
                    server_type=first_key,
                    name=second_key,
                    dir_path=str(dir_path),
                    entrypoint=str(entrypoint_fpath),
                    host=host,
                    port=port,
                    url=f"http://{host}:{port}" if host and port else None,
                    pid=process.pid,
                    config_path=top_level_path,
                    start_time=start_time,
                )
            )

        self._head_server_instance.set_server_instances(
            [inst.model_dump(mode="json") for inst in self._server_instance_display_configs]
        )

        self._server_client = ServerClient(
            head_server_config=ServerClient.load_head_server_config(),
            global_config_dict=global_config_dict,
        )

        print("Waiting for head server to spin up")
        poll_count = 0
        while True:
            status = self._server_client.poll_for_status(HEAD_SERVER_KEY_NAME)
            if status == "success":
                break

            if poll_count % 10 == 0:  # Print every 30s
                print(f"Head server is not up yet (status `{status}`). Sleeping...")

            poll_count += 1
            sleep(3)

        print("Waiting for servers to spin up")
        if global_config_dict[DRY_RUN_KEY_NAME]:
            self.wait_for_dry_run_spinup()
        else:
            self.wait_for_spinup()

    def display_server_instance_info(self) -> None:
        if not self._server_instance_display_configs:
            print("No server instances to display.")
            return

        print(f"""
{"#" * 100}
#
# Server Instances
#
{"#" * 100}
""")

        for i, inst in enumerate(self._server_instance_display_configs, 1):
            print(f"[{i}] {inst.process_name} ({inst.server_type}/{inst.name})")
            pprint(inst.model_dump(mode="json", exclude={"start_time", "status", "uptime_seconds"}))
        print(f"{'#' * 100}\n")

    def poll(self) -> None:
        if not self._head_server_thread.is_alive():
            raise RuntimeError("Head server finished unexpectedly!")

        for process_name, process in self._processes.items():
            if process.poll() is not None:
                proc_out, proc_err = process.communicate()
                print_str = f"Process `{process_name}` finished unexpectedly!"

                if isinstance(proc_out, bytes):
                    proc_out = proc_out.decode("utf-8")
                    print_str = f"""{print_str}
Process `{process_name}` stdout:
{proc_out}
"""
                if isinstance(proc_err, bytes):
                    proc_err = proc_err.decode("utf-8")
                    print_str = f"""{print_str}
Process `{process_name}` stderr:
{proc_err}"""

                raise RuntimeError(print_str)

    def wait_for_dry_run_spinup(self) -> None:
        sleep_interval = 3

        remaining_processes = list(self._processes.values())
        while remaining_processes:
            for i in reversed(range(len(remaining_processes))):
                process = remaining_processes[i]
                if process.poll() is not None:
                    remaining_processes.pop(i)

            sleep(sleep_interval)

    def wait_for_spinup(self) -> None:
        sleep_interval = 3
        poll_count = 0
        successful_servers = []
        total_servers = len(self._server_instance_display_configs)

        # Until we spin up or error out.
        while True:
            self.poll()
            statuses = self.check_http_server_statuses(successful_servers)
            successful_servers.extend(s for s, status in statuses if status == "success")

            waiting = []
            for name, status in statuses:
                if status != "success":
                    waiting.append(name)

            if len(successful_servers) != total_servers:
                if poll_count % 10 == 0:  # Print every sleep_interval * poll_count = 3 * 10 = 30s
                    print(
                        f"""Checking for HTTP server statuses.
{len(successful_servers)} / {total_servers} servers ready. Waiting for servers to spin up: {waiting}"""
                    )
                poll_count += 1
            else:
                print(f"All {len(successful_servers)} / {total_servers} servers ready! Polling every 60s")
                self.display_server_instance_info()
                return

            sleep(sleep_interval)

    def shutdown(self) -> None:
        print("Sending interrupt signals to servers...")
        for process in self._processes.values():
            process.send_signal(SIGINT)

        print("Waiting for processes to finish...")
        killed_process_names: List[str] = []
        unreaped_process_names: List[str] = []
        for process_name, process in self._processes.items():
            try:
                process.wait(timeout=_GRACEFUL_SHUTDOWN_TIMEOUT_SEC)
            except TimeoutExpired:
                process.kill()
                killed_process_names.append(process_name)
                # Reap the child after SIGKILL to avoid leaving a <defunct> entry.
                try:
                    process.wait(timeout=_FORCE_KILL_REAP_TIMEOUT_SEC)
                except TimeoutExpired:
                    unreaped_process_names.append(process_name)

        if killed_process_names:
            print(
                f"""Some processes ({", ".join(killed_process_names)}) didn't shutdown within the {_GRACEFUL_SHUTDOWN_TIMEOUT_SEC}s timeout, killing instead. You may see messages like:
```bash
rpc_client.h:203: Failed to connect to GCS within 60 seconds. GCS may have been killed. It's either GCS is terminated by `ray stop` or is killed unexpectedly. If it is killed unexpectedly, see the log file gcs_server.out. https://docs.ray.io/en/master/ray-observability/user-guides/configure-logging.html#logging-directory-structure. The program will terminate.
```
"""
            )
        if unreaped_process_names:
            print(
                f"WARNING: processes ({', '.join(unreaped_process_names)}) did not exit "
                f"within {_FORCE_KILL_REAP_TIMEOUT_SEC}s after SIGKILL; "
                "they may remain as zombies until this process exits."
            )
        self._processes = dict()

        self._head_server.should_exit = True
        self._head_server_thread.join()

        self._head_server = None
        self._head_server_thread = None

        print("NeMo Gym finished!")

    def run_forever(self) -> None:
        if self._server_client.global_config_dict[DRY_RUN_KEY_NAME]:
            self.shutdown()
            return

        async def sleep():
            # Indefinitely
            while True:
                self.poll()
                await asyncio.sleep(60)

        try:
            asyncio.run(sleep())
        except KeyboardInterrupt:
            pass
        finally:
            self.shutdown()

    def check_http_server_statuses(self, successful_servers: List[str]) -> List[Tuple[str, ServerStatus]]:
        statuses = []
        for server_instance_display_config in self._server_instance_display_configs:
            name = server_instance_display_config.config_path

            # No need to re-poll successfully spun up servers.
            if name in successful_servers:
                continue

            status = self._server_client.poll_for_status(name)
            statuses.append((name, status))

        return statuses


@exit_cleanly_on_config_error
def run(
    global_config_dict_parser_config: Optional[GlobalConfigDictParserConfig] = None,
):  # pragma: no cover
    """
    Start NeMo Gym servers for agents, models, and resources.

    This command reads configuration from YAML files specified via `+config_paths` and starts all configured servers.
    The configuration files should define server instances with their entrypoints and settings.

    Configuration Parameter:
        config_paths (List[str]): Paths to YAML configuration files. Specify via Hydra: `+config_paths="[file1.yaml,file2.yaml]"`

    Examples:

    ```bash
    # Start servers with specific configs
    config_paths="resources_servers/example_single_tool_call/configs/example_single_tool_call.yaml,\\
    responses_api_models/openai_model/configs/openai_model.yaml"
    gym env start "+config_paths=[${config_paths}]"
    ```
    """
    global_config_dict = get_global_config_dict(global_config_dict_parser_config=global_config_dict_parser_config)
    # Just here for help
    BaseNeMoGymCLIConfig.model_validate(global_config_dict)

    rh = RunHelper()
    rh.start(global_config_dict_parser_config)
    rh.run_forever()


def _validate_data_single(test_config: TestConfig) -> None:  # pragma: no cover
    if not test_config.should_validate_data:
        return

    # We have special data checks for resources servers
    if test_config.dir_path.parts[0] != "resources_servers":
        return

    # Check that the required examples and example metrics are present. Read from the resolved dir
    # (built-ins live under the install root) while messages reference the relative entrypoint.
    example_fpath = test_config.resolved_dir_path / "data/example.jsonl"
    assert example_fpath.exists(), (
        f"A jsonl file containing 5 examples is required for the {test_config.dir_path} resources server. The file must be found at {example_fpath}. Usually this example data is just the first 5 examples of your train dataset."
    )
    with open(example_fpath) as f:
        count = sum(1 for _ in f)
    assert count == 5, f"Expected 5 examples at {example_fpath} but got {count}."

    server_type_name = test_config.dir_path.parts[-1]
    example_metrics_fpath = test_config.resolved_dir_path / "data/example_metrics.json"
    assert (
        example_metrics_fpath.exists()
    ), f"""You must run the example data validation for the example data found at {example_fpath}.
Your command should look something like the following (you should update this command with your actual server config path):
```bash
gym dataset collate "+config_paths=[{test_config._dir_path}/configs/{server_type_name}.yaml]" \\
    +output_dirpath=data/{server_type_name} \\
    +mode=example_validation
```
and your config must include an agent server config with an example dataset like:
```yaml
example_multi_step_simple_agent:
  responses_api_agents:
    simple_agent:
      ...
      datasets:
      - name: example
        type: example
        jsonl_fpath: resources_servers/example_multi_step/data/example.jsonl
```

See `resources_servers/example_multi_step/configs/example_multi_step.yaml` for a full config example.
"""
    with open(example_metrics_fpath) as f:
        example_metrics = json.load(f)
    assert example_metrics["Number of examples"] == 5, (
        f"Expected 5 examples in the metrics at {example_metrics_fpath}, but got {example_metrics['Number of examples']}"
    )

    conflict_paths = glob(str(test_config.resolved_dir_path / "data/*conflict*"))
    conflict_paths_str = "\n- ".join([""] + conflict_paths)
    assert not conflict_paths, f"Found {len(conflict_paths)} conflicting paths: {conflict_paths_str}"

    example_rollouts_fpath = test_config.resolved_dir_path / "data/example_rollouts.jsonl"
    assert example_rollouts_fpath.exists(), f"""You must run the example data through your agent and provide the example rollouts at `{example_rollouts_fpath}`.

Your commands should look something like:
```bash
# Server spinup
example_multi_step_config_paths="responses_api_models/openai_model/configs/openai_model.yaml,\
resources_servers/example_multi_step/configs/example_multi_step.yaml"
gym env start "+config_paths=[${{example_multi_step_config_paths}}]"

# Collect the rollouts
gym eval run --no-serve +agent_name=example_multi_step_simple_agent \
    +input_jsonl_fpath=resources_servers/example_multi_step/data/example.jsonl \
    +output_jsonl_fpath=resources_servers/example_multi_step/data/example_rollouts.jsonl \
    +limit=null

# View your rollouts
head -1 resources_servers/example_multi_step/data/example_rollouts.jsonl
```
"""
    with open(example_rollouts_fpath) as f:
        count = sum(1 for _ in f)
    assert count == 5, f"Expected 5 example rollouts in {example_rollouts_fpath}, but got {count}"

    print(f"The data for {test_config.dir_path} has been successfully validated!")


def _test_single(test_config: TestConfig, global_config_dict: DictConfig) -> Popen:  # pragma: no cover
    # Eventually we may want more sophisticated testing here, but this is sufficient for now.
    prefix = test_config.entrypoint.replace("/", "\\/")
    resolved_dir = test_config.resolved_dir_path
    command = f"""{setup_env_command(resolved_dir, global_config_dict, prefix)} && pytest"""
    # Generated server tests import `resources_servers.<name>...`, so the project root (the dir
    # holding the server-type dirs) must be on PYTHONPATH when running from outside a repo checkout.
    return run_command(command, resolved_dir, project_root=resolved_dir.parent.parent)


def test():  # pragma: no cover
    global_config_dict = get_global_config_dict()
    test_config = TestConfig.model_validate(global_config_dict)

    proc = _test_single(test_config, global_config_dict)
    return_code = proc.wait()
    if return_code != 0:
        print(f"You can run detailed tests via `cd {test_config.entrypoint} && source .venv/bin/activate && pytest`.")
        exit(return_code)

    try:
        _validate_data_single(test_config)
    except AssertionError:
        print(f"Data validation failed for {test_config.entrypoint}. You can rerun just the data validation like:")
        print("```bash")
        print(f"gym env test +entrypoint={test_config.entrypoint} +should_validate_data=true")
        print("```")
        exit(1)


def _display_list_of_paths(paths: List[Path]) -> str:  # pragma: no cover
    paths = list(map(str, paths))
    return "".join(f"\n- {p}" for p in paths)


def _format_pct(count: int, total: int) -> str:  # pragma: no cover
    return f"{count} / {total} ({100 * count / total:.2f}%)"


class TestAllConfig(BaseNeMoGymCLIConfig):
    """
    Run tests for all server modules in the project.

    Examples:

    ```bash
    gym env test
    ```
    """

    fail_on_total_and_test_mismatch: bool = Field(
        default=False,
        description="Fail if the number of server modules doesn't match the number with tests (default: False).",
    )
    delete_venvs_after_each_test: bool = Field(
        default=False,
        description="Delete each server venv after its tests have been run (default: False).",
    )
    num_shards: int = Field(
        default=1,
        ge=1,
        description="Total number of shards to split the server suite across (default: 1 = no sharding). "
        "Used to parallelize the suite across CI runners.",
    )
    shard_index: int = Field(
        default=0,
        ge=0,
        description="Which shard (0-based) this invocation runs; must be < num_shards (default: 0).",
    )


def _select_shard(dir_paths: List[Path], shard_index: int, num_shards: int) -> List[Path]:
    """Deterministically select this shard's subset of modules.

    Round-robin (stride) over a sorted list spreads heavy modules across shards more evenly than
    contiguous chunks, which balances wall-time when the suite is parallelized across CI runners.
    """
    if num_shards <= 1:
        return dir_paths
    assert 0 <= shard_index < num_shards, (
        f"shard_index ({shard_index}) must be in [0, num_shards) for num_shards={num_shards}"
    )
    return sorted(dir_paths, key=str)[shard_index::num_shards]


def test_all():  # pragma: no cover
    global_config_dict = get_global_config_dict()
    test_all_config = TestAllConfig.model_validate(global_config_dict)

    # Discover server modules under both the cwd (a user's project) and the Gym install root
    # (built-ins, which live under PARENT_DIR in editable and wheel installs). Entrypoints are kept
    # relative; the cwd shadows the install root for same-named modules. This lets `gym env test`
    # discover and run built-in servers from any cwd, not only a repo checkout.
    server_type_dirs = ("resources_servers", "responses_api_agents", "responses_api_models")
    seen_rel_paths: set[str] = set()
    candidate_dir_paths: List[str] = []
    for root in (Path.cwd(), PARENT_DIR):
        for server_type_dir in server_type_dirs:
            for module_path in sorted((root / server_type_dir).glob("*")):
                if "pycache" in module_path.name or not module_path.is_dir():
                    continue
                rel_path = f"{server_type_dir}/{module_path.name}"
                if rel_path in seen_rel_paths:
                    continue
                seen_rel_paths.add(rel_path)
                candidate_dir_paths.append(rel_path)
    print(f"Found {len(candidate_dir_paths)} total modules:{_display_list_of_paths(candidate_dir_paths)}\n")
    dir_paths: List[Path] = list(map(Path, candidate_dir_paths))
    dir_paths = [p for p in dir_paths if (_resolve_server_dir(p) / "README.md").exists()]
    print(f"Found {len(dir_paths)} modules to test:{_display_list_of_paths(dir_paths)}\n")

    # Keep the full list for the total-vs-tested mismatch check below, then narrow to this shard.
    full_dir_paths = dir_paths
    dir_paths = _select_shard(dir_paths, test_all_config.shard_index, test_all_config.num_shards)
    if test_all_config.num_shards > 1:
        print(
            f"Shard {test_all_config.shard_index + 1}/{test_all_config.num_shards}: "
            f"testing {len(dir_paths)} of {len(full_dir_paths)} modules:{_display_list_of_paths(dir_paths)}\n"
        )

    tests_passed: List[Path] = []
    tests_failed: List[Path] = []
    tests_missing: List[Path] = []
    tests_unrecognized: List[Path] = []
    data_validation_failed: List[Path] = []
    times_taken: List[Tuple[float, Path]] = []
    for dir_path in tqdm(dir_paths, desc="Running tests"):
        start_time = time()

        test_config = TestConfig(
            entrypoint=str(dir_path),
            should_validate_data=True,  # Test all always validates data.
        )
        proc = _test_single(test_config, global_config_dict)
        return_code = proc.wait()

        match return_code:
            case 0:
                tests_passed.append(dir_path)
            case 1 | 2:
                tests_failed.append(dir_path)
            case 5:
                tests_missing.append(dir_path)
            case _:
                tests_unrecognized.append(dir_path)

        try:
            _validate_data_single(test_config)
        except AssertionError:
            data_validation_failed.append(dir_path)

        if test_all_config.delete_venvs_after_each_test:
            venv_path = _resolve_server_dir(dir_path) / ".venv"
            print(f"Deleting {venv_path} since `delete_venvs_after_each_test=true`")
            rmtree(venv_path, ignore_errors=True)

        times_taken.append((time() - start_time, dir_path))

    times_taken.sort(reverse=True)
    table = Table(title="Times taken per test (sorted from highest to lowest)")
    table.add_column("Server path")
    table.add_column("Time taken (s)")
    for time_taken, dir_path in times_taken:
        table.add_row(str(dir_path), f"{time_taken:.2f}")
    rich.print(table)

    print(f"""Found {len(candidate_dir_paths)} total modules:{_display_list_of_paths(candidate_dir_paths)}

Found {len(dir_paths)} modules to test:{_display_list_of_paths(dir_paths)}

Tests passed {_format_pct(len(tests_passed), len(dir_paths))}:{_display_list_of_paths(tests_passed)}

Tests failed {_format_pct(len(tests_failed), len(dir_paths))}:{_display_list_of_paths(tests_failed)}

Tests missing {_format_pct(len(tests_missing), len(dir_paths))}:{_display_list_of_paths(tests_missing)}

Tests that returned unrecognized exit codes {_format_pct(len(tests_unrecognized), len(dir_paths))}:{_display_list_of_paths(tests_unrecognized)}

Data validation failed {_format_pct(len(data_validation_failed), len(dir_paths))}:{_display_list_of_paths(data_validation_failed)}
""")

    if tests_failed or tests_missing or tests_unrecognized:
        print(f"""You can rerun just the server with failed, missing or unrecognized tests results like:
```bash
gym env test +entrypoint={(tests_failed + tests_missing + tests_unrecognized)[0]}
```
""")
    if data_validation_failed:
        print(f"""You can rerun just the server with failed data validation like:
```bash
gym env test +entrypoint={data_validation_failed[0]} +should_validate_data=true
```
""")

    if test_all_config.fail_on_total_and_test_mismatch:
        # Compare against the full (unsharded) module list — every module must be testable
        # regardless of how many shards we split the run into.
        extra_candidates = [p for p in candidate_dir_paths if Path(p) not in full_dir_paths]
        assert (
            len(candidate_dir_paths) == len(full_dir_paths)
        ), f"""Mismatch on the number of total modules found ({len(candidate_dir_paths)}) and the number of actual modules tested ({len(full_dir_paths)})!

Extra candidate paths:{_display_list_of_paths(extra_candidates)}"""

    if tests_missing or tests_failed or data_validation_failed:
        exit(1)


def init_resources_server():  # pragma: no cover
    """
    Initialize a new resources server with template files and directory structure.

    Examples:

    ```bash
    gym env init --resources-server my_server
    ```
    """
    config_dict = get_global_config_dict()
    run_config = RunConfig.model_validate(config_dict)

    dirpath = Path(run_config.entrypoint).resolve()

    if exists(dirpath):
        print(f"Folder already exists: {dirpath}. Exiting init.")
        exit()

    makedirs(dirpath)

    server_type = "resources_servers"
    server_type_name = dirpath.parts[-1].lower()
    server_type_title = "".join(x.capitalize() for x in server_type_name.split("_"))

    configs_dirpath = dirpath / "configs"
    makedirs(configs_dirpath)

    config_fpath = configs_dirpath / f"{server_type_name}.yaml"
    with open(config_fpath, "w") as f:
        f.write(f"""# Resources server: owns this environment's task verification (verify()) and reward.
{server_type_name}_resources_server:          # instance name — how agents/CLI refer to this server
  {server_type}:                              # server type: resources_servers | responses_api_agents | responses_api_models
    {server_type_name}:                        # implementation directory under {server_type}/
      entrypoint: app.py                        # server entry module
      domain: other                             # task domain; one of: math, coding, agent, knowledge,
                                                #   instruction_following, long_context, safety, games, translation,
                                                #   e2e, rlhf, other. Change 'other' to the closest fit.
      verified: false                           # set true once the benchmark has been baselined and reviewed

# Agent server config specifies the agent server to run and any additional components of the environment such as resources servers
{server_type_name}_simple_agent:               # this agent instance's name — pass as --agent to gym eval run
  responses_api_agents:
    simple_agent:                               # built-in agent: runs the model with tool calls (up to max_steps); swap for your own agent dir
      entrypoint: app.py
      resources_server:                         # the resources server this agent interacts with for tools, state and verification
        type: resources_servers
        name: {server_type_name}_resources_server
      model_server:                             # the model that answers; 'policy_model' is resolved from a model config
        type: responses_api_models
        name: policy_model
      datasets:                                 # one block per split: train | validation | example
      - name: train
        type: train
        jsonl_fpath: resources_servers/{server_type_name}/data/train.jsonl   # local data file for this split
        num_repeats: 1                          # times to repeat each example (e.g. for pass@k / mean@k)
        license: Apache 2.0                     # required for train/validation; must be an allowed license string
        # to fetch this split from a registry instead, add a source: block (type: gitlab | huggingface)
      - name: validation
        type: validation
        jsonl_fpath: resources_servers/{server_type_name}/data/validation.jsonl
        num_repeats: 1
        license: Apache 2.0
      - name: example                           # 5 rows committed to git for quick smoke tests
        type: example
        jsonl_fpath: resources_servers/{server_type_name}/data/example.jsonl
        num_repeats: 1
""")

    app_fpath = dirpath / "app.py"
    with open(ROOT_DIR / "resources/resources_server_template.py") as f:
        app_template = f.read()
    app_content = app_template.replace("ExampleMultiStep", server_type_title)
    with open(app_fpath, "w") as f:
        f.write(app_content)

    tests_dirpath = dirpath / "tests"
    makedirs(tests_dirpath)

    tests_fpath = tests_dirpath / "test_app.py"
    with open(ROOT_DIR / "resources/resources_server_test_template.py") as f:
        tests_template = f.read()
    tests_content = tests_template.replace("ExampleMultiStep", server_type_title)
    tests_content = tests_content.replace("from app", f"from resources_servers.{server_type_name}.app")
    with open(tests_fpath, "w") as f:
        f.write(tests_content)

    requirements_fpath = dirpath / "requirements.txt"
    with open(requirements_fpath, "w") as f:
        if (PARENT_DIR / "pyproject.toml").exists():
            # local nemo gym - detected by ../pyproject.toml exists
            rel_to_gym_root = os.path.relpath(PARENT_DIR, dirpath)
            f.write(f"-e nemo-gym[dev] @ {rel_to_gym_root}\n")
        else:
            # pypi path
            f.write("nemo-gym[dev]\n")

    readme_fpath = dirpath / "README.md"
    with open(readme_fpath, "w") as f:
        f.write("""# Description

Data links: ?

# Licensing information
Code: ?
Data: ?

Dependencies
- nemo_gym: Apache 2.0
?
""")

    data_dirpath = dirpath / "data"
    data_dirpath.mkdir(exist_ok=True)

    data_gitignore_fpath = data_dirpath / ".gitignore"
    with open(data_gitignore_fpath, "w") as f:
        f.write("""*train.jsonl
*validation.jsonl
*train_prepare.jsonl
*validation_prepare.jsonl
*example_prepare.jsonl
""")


def dump_config():  # pragma: no cover
    """
    Display the resolved Hydra configuration for debugging purposes.

    Examples:

    ```bash
    gym env resolve "+config_paths=[<config1>,<config2>]"
    ```
    """
    global_config_dict = get_global_config_dict(
        global_config_dict_parser_config=GlobalConfigDictParserConfig(
            hide_secrets=True,
        ),
    )

    # Just here for help
    BaseNeMoGymCLIConfig.model_validate(global_config_dict)

    print(OmegaConf.to_yaml(global_config_dict, resolve=True))


@exit_cleanly_on_config_error
def validate():
    """Validate a config without starting Ray or any server subprocess.

    Runs the full config parse — config_paths resolution (missing/malformed), server cross-reference
    validation, mandatory `???` values, and schema — then exits 0 (valid) or, via
    `exit_cleanly_on_config_error`, 1 with a clean traceback-free message. No Ray, no servers, so it
    returns in well under a second instead of after Ray bootstrap.

    No model config is required: a dummy `policy_model` is injected (the `NO_MODEL` parser config, as
    in `gym list` / `env compose`) so model interpolations (e.g. `${policy_base_url}`) resolve —
    validation is about config well-formedness, not the model. Pass a model config / `--model-type`
    as well if you want it validated too.

    Examples:

    ```bash
    gym env validate --environment <env>
    gym env validate --benchmark <benchmark>
    # or by explicit config path(s):
    gym env validate --config resources_servers/<env>/configs/<env>.yaml
    ```
    """
    global_config_dict = get_global_config_dict(
        global_config_dict_parser_config=GlobalConfigDictParserConfig(
            initial_global_config_dict=GlobalConfigDictParserConfig.NO_MODEL_GLOBAL_CONFIG_DICT,
        ),
    )
    BaseNeMoGymCLIConfig.model_validate(global_config_dict)

    rich.print("[green]✓[/green] Config is valid.")


def list_environments() -> None:
    """List the environments available under environments/, by short name.

    Examples:

    ```bash
    gym list environments
    gym list environments --json
    ```
    """
    global_config_dict = get_global_config_dict(
        global_config_dict_parser_config=GlobalConfigDictParserConfig(
            initial_global_config_dict=GlobalConfigDictParserConfig.NO_MODEL_GLOBAL_CONFIG_DICT,
        )
    )
    BaseNeMoGymCLIConfig.model_validate(global_config_dict)

    environments = discover_environments()

    if global_config_dict.get(JSON_OUTPUT_KEY_NAME, False):
        print(
            json.dumps(
                [
                    {"name": name, "domain": env.domain, "description": env.description}
                    for name, env in environments.items()
                ]
            )
        )
        return

    if not environments:
        rich.print("[yellow]No environments found.[/yellow]")
        return

    table = Table(title=f"Available environments in NeMo Gym ({len(environments)})")
    table.add_column("Environment")
    table.add_column("Domain")
    table.add_column("Description")
    for name, environment in environments.items():
        table.add_row(name, environment.domain or "", environment.description or "")

    rich.print(table)


def status():  # pragma: no cover
    global_config_dict = get_global_config_dict()
    BaseNeMoGymCLIConfig.model_validate(global_config_dict)

    status_cmd = StatusCommand()
    servers = status_cmd.discover_servers()

    if global_config_dict.get(JSON_OUTPUT_KEY_NAME, False):
        print(json.dumps([server.model_dump(mode="json") for server in servers]))
        return

    status_cmd.display_status(servers)


class PipListConfig(RunConfig):
    format: Optional[str] = Field(
        default=None,
        description="Output format for pip list. Options: 'columns' (default), 'freeze', 'json'",
    )
    outdated: bool = Field(
        default=False,
        description="List outdated packages",
    )


def pip_list():  # pragma: no cover
    """List packages installed in a server's virtual environment."""
    global_config_dict = get_global_config_dict()
    config = PipListConfig.model_validate(global_config_dict)

    dir_path = _resolve_server_dir(Path(config.entrypoint))
    venv_path = dir_path / ".venv"

    if not venv_path.exists():
        print(f"  Virtual environment not found at: {venv_path}")
        print("  Run tests or setup the server first using:")
        print(f"  gym env test +entrypoint={config.entrypoint}")
        exit(1)

    pip_list_cmd = "uv pip list"
    if config.format:
        pip_list_cmd += f" --format={config.format}"
    if config.outdated:
        pip_list_cmd += " --outdated"

    command = f"""cd {dir_path} \\
    && source .venv/bin/activate \\
    && {pip_list_cmd}"""

    print(f"  Package list for: {config.entrypoint}")
    print(f"Virtual environment: {venv_path.absolute()}")
    print("-" * 72)

    proc = run_command(command, dir_path)
    return_code = proc.wait()
    exit(return_code)
