# SPDX-FileCopyrightText: Copyright (c) 2025 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 warnings
from argparse import ArgumentParser
from enum import Enum
from pathlib import Path
from typing import Annotated, Any, ClassVar, Dict, List, Literal, Optional, Set, Tuple, Union

import rich
from omegaconf import DictConfig, OmegaConf
from pydantic import (
    BaseModel,
    ConfigDict,
    Field,
    TypeAdapter,
    ValidationError,
    model_validator,
)
from pydantic_core import PydanticUndefined
from rich.markdown import Markdown
from rich.text import Text


########################################
# Base CLI configs
########################################


class BaseNeMoGymCLIConfig(BaseModel):
    @model_validator(mode="before")
    @classmethod
    def pre_process(cls, data):
        parser = ArgumentParser(add_help=False)
        parser.add_argument("-h", "--help", action="store_true")
        args, _ = parser.parse_known_args()

        if not (args.help or data.get("h") or data.get("help")):
            return data

        rich.print(f"""Displaying help for [bold]{cls.__name__}[/bold]
""")
        # We use __doc__ directly here since inspect.getdoc will inherit the doc from parent classes.
        class_doc = cls.__doc__
        if class_doc:
            rich.print(f"""[bold]Description[/bold]
-----------
{class_doc.strip()}
""")
            # Render docstring as Markdown
            md = Markdown(class_doc.strip())
            rich.print(md)

        fields = cls.model_fields.items()
        if fields:
            rich.print("""[bold]Parameters[/bold]
----------""")

            prefixes: List[Text] = []
            suffixes: List[Text] = []
            for field_name, field in fields:
                description_str = field.description if field.description else ""

                # Not sure if there is a better way to get this annotation_str, e.g. using typing.get_args or typing.get_origin
                annotation_str = (
                    field.annotation.__name__ if isinstance(field.annotation, type) else str(field.annotation)
                )
                annotation_str = annotation_str.replace("typing.", "")

                # Add default value information if available
                if field.default is not PydanticUndefined and field.default is not None:
                    default_str = f" [default: {field.default}]"
                    description_str = description_str + default_str if description_str else default_str.strip()
                elif field.default_factory is not None:
                    default_str = " [default: <factory>]"
                    description_str = description_str + default_str if description_str else default_str.strip()
                elif field.default is PydanticUndefined and field.is_required():
                    default_str = " [required]"
                    description_str = description_str + default_str if description_str else default_str.strip()

                prefixes.append(Text.from_markup(f"- [blue]{field_name}[/blue] [yellow]({annotation_str})[/yellow]"))
                suffixes.append(description_str)

            max_prefix_length = max(map(len, prefixes))
            ljust_length = max_prefix_length + 3
            for prefix, suffix in zip(prefixes, suffixes):
                prefix.align("left", ljust_length)
                rich.print(prefix + suffix)
        else:
            print("There are no arguments to this CLI command!")

        # Exit after help is printed.
        exit()


########################################
# Server references
#
# We enable servers to reference other servers. The way they do so is through these schemas below.
########################################


class ModelServerRef(BaseModel):
    type: Literal["responses_api_models"]
    name: str


class ResourcesServerRef(BaseModel):
    type: Literal["resources_servers"]
    name: str


class AgentServerRef(BaseModel):
    type: Literal["responses_api_agents"]
    name: str


ServerRef = Union[ModelServerRef, ResourcesServerRef, AgentServerRef]
ServerRefTypeAdapter = TypeAdapter(ServerRef)


def is_server_ref(config_dict: DictConfig) -> Optional[ServerRef]:
    try:
        return ServerRefTypeAdapter.validate_python(config_dict)
    except ValidationError:
        return None


class ConfigError(Exception):
    """Base for user-facing configuration errors.

    These represent actionable user mistakes (typos, missing files, malformed input) rather than
    internal bugs. The CLI catches `ConfigError` and prints just the message — no traceback —
    while still leaving them as ordinary exceptions so callers like `validate` can catch and
    format them.
    """


class ConfigPathNotFoundError(ConfigError, FileNotFoundError):
    """A `config_paths` entry could not be found in the cwd or the Gym install location."""


class MalformedConfigPathsError(ConfigError, ValueError):
    """`config_paths` was not a list of paths (e.g. a scalar string was passed)."""


class NoServerInstancesError(ConfigError, ValueError):
    """A run was requested but the merged config defines no server instances to start."""


class ConfigMissingValuesError(ConfigError, ValueError):
    """One or more required config values are still unset (OmegaConf '???') after merging."""


class ServerRefNotFoundError(ConfigError, ValueError):
    """A server cross-reference points to an instance that is not defined in the merged config."""


class InheritPathNotFoundError(ConfigError, ValueError):
    """An `_inherit_from` / swap / copy directive references a config path that does not exist."""


class AlmostServerError(ConfigError, ValueError):
    """One or more server blocks are almost-servers (right shape, failed validation) and
    `error_on_almost_servers` is set, so the run is aborted."""


########################################
# Dataset configs for handling and upload/download
########################################


class UploadJsonlDatasetGitlabConfig(BaseNeMoGymCLIConfig):
    """
    Upload a local jsonl dataset artifact to Gitlab.

    Examples:

    ```bash
    gym dataset upload --storage gitlab \
        +dataset_name=example_multi_step \
        +version=0.0.1 \
        +input_jsonl_fpath=data/train.jsonl
    ```
    """

    dataset_name: str = Field(description="The dataset name.")
    version: str = Field(description="The version of this dataset. Must be in the format `x.x.x`.")
    input_jsonl_fpath: str = Field(description="Path to the jsonl file to upload.")


class JsonlDatasetGitlabIdentifer(BaseModel):
    dataset_name: str
    version: str
    artifact_fpath: str


class DownloadJsonlDatasetGitlabConfig(JsonlDatasetGitlabIdentifer, BaseNeMoGymCLIConfig):
    """
    Download a JSONL dataset from GitLab Model Registry.

    Examples:

    ```bash
    gym dataset download --storage gitlab \
        +dataset_name=example_multi_step \
        +version=0.0.1 \
        +artifact_fpath=train.jsonl \
        +output_fpath=data/train.jsonl
    ```
    """

    dataset_name: str = Field(description="The dataset name.")
    version: str = Field(description="The version of this dataset. Must be in the format `x.x.x`.")
    artifact_fpath: str = Field(description="The filepath to the artifact to download.")
    output_fpath: str = Field(description="Where to save the downloaded dataset.")


class DeleteJsonlDatasetGitlabConfig(BaseNeMoGymCLIConfig):
    """
    Delete a dataset from GitLab Model Registry (prompts for confirmation).

    Examples:

    ```bash
    gym dataset rm +dataset_name=old_dataset
    ```
    """

    dataset_name: str = Field(description="Name of the dataset to delete from GitLab.")


class JsonlDatasetHuggingFaceIdentifer(BaseModel):
    repo_id: str = Field(description="The repo id.")
    artifact_fpath: Optional[str] = Field(
        default=None,
        description="Path to specific file in HuggingFace repo (e.g., 'train.jsonl'). If omitted, load_dataset will be used with split.",
    )


class BaseUploadJsonlDatasetHuggingFaceConfig(BaseNeMoGymCLIConfig):
    """
    Upload a JSONL dataset to HuggingFace Hub with automatic naming based on domain and resources server.

    Examples:

    ```bash
    resource_config_path="resources_servers/example_multi_step/configs/example_multi_step.yaml"
    gym dataset upload \
        +dataset_name=my_dataset \
        +input_jsonl_fpath=data/train.jsonl \
        +resource_config_path=${resource_config_path}
    ```
    """

    # Must match `nemo_gym.global_config.HF_TOKEN_KEY_NAME`
    hf_token: str = Field(description="HuggingFace API token for authentication.")
    hf_organization: str = Field(description="HuggingFace organization name where dataset will be uploaded.")
    hf_collection_name: str = Field(description="HuggingFace collection name for organizing datasets.")
    hf_collection_slug: str = Field(description="Alphanumeric collection slug found at the end of collection URI.")
    dataset_name: Optional[str] = Field(
        default=None, description="Name of the dataset (will be combined with domain and resources server name)."
    )
    input_jsonl_fpath: str = Field(description="Path to the local jsonl file to upload.")
    resource_config_path: str = Field(
        description="Path to resources server config file (used to extract domain for naming convention)."
    )
    hf_dataset_prefix: str = Field(
        default="Nemotron-RL", description="Prefix prepended to dataset name (default: 'NeMo-Gym')."
    )
    split: Literal["train", "validation", "test"] = Field(
        default="train",
        description="Dataset split type (e.g., 'train', 'validation', 'test'). Format validation only applies to 'train' splits.",
    )
    create_pr: bool = Field(
        default=False,
        description="Create a pull request instead of pushing directly. Required for repos where you do not have write access.",
    )
    revision: Optional[str] = Field(
        default=None,
        description="Git revision (branch name) to upload to. Use the same revision for multiple files to upload to the same PR. If not provided with create_pr=True, a new branch/PR will be created automatically.",
    )
    commit_message: Optional[str] = Field(
        default=None, description="Custom commit message. If not provided, HuggingFace auto-generates one."
    )
    commit_description: Optional[str] = Field(
        default=None, description="Optional commit description with additional context."
    )


class UploadJsonlDatasetHuggingFaceConfig(BaseUploadJsonlDatasetHuggingFaceConfig):
    """
    Upload a JSONL dataset to HuggingFace Hub and automatically delete from GitLab after successful upload.

    This command always deletes the dataset from GitLab after uploading to HuggingFace.
    Use `gym dataset upload` if you want optional deletion control.

    Examples:

    ```bash
    resource_config_path="resources_servers/example_multi_step/configs/example_multi_step.yaml"
    gym dataset migrate \
        +dataset_name=my_dataset \
        +input_jsonl_fpath=data/train.jsonl \
        +resource_config_path=${resource_config_path}
    ```
    """

    forbidden_fields: ClassVar[Set[str]] = {"delete_from_gitlab"}

    @model_validator(mode="before")
    def check_forbidden_fields(cls, data):
        if isinstance(data, dict) or hasattr(data, "keys"):
            forbidden = cls.forbidden_fields.intersection(set(data.keys()))
            if forbidden:
                raise ValueError(f"Forbidden fields present: {forbidden}")
        return data


class UploadJsonlDatasetHuggingFaceMaybeDeleteConfig(BaseUploadJsonlDatasetHuggingFaceConfig):
    """
    Upload a JSONL dataset to HuggingFace Hub with optional GitLab deletion after successful upload.

    Examples:

    ```bash
    resource_config_path="resources_servers/example_multi_step/configs/example_multi_step.yaml"
    gym dataset upload \
        +dataset_name=my_dataset \
        +input_jsonl_fpath=data/train.jsonl \
        +resource_config_path=${resource_config_path} \
        +delete_from_gitlab=true
    ```
    """

    delete_from_gitlab: Optional[bool] = Field(
        default=False, description="Delete the dataset from GitLab after successful upload to HuggingFace."
    )


class DownloadJsonlDatasetHuggingFaceConfig(JsonlDatasetHuggingFaceIdentifer, BaseNeMoGymCLIConfig):
    """
    Download a JSONL dataset from HuggingFace Hub to local filesystem.

    Examples:

    ```bash
    gym dataset download \
        +repo_id=NVIDIA/NeMo-Gym-Math-example_multi_step-v1 \
        +artifact_fpath=train.jsonl \
        +output_fpath=data/train.jsonl
    ```
    """

    output_dirpath: Optional[str] = Field(
        default=None,
        description="Directory to save the downloaded dataset. Files will be named {split}.jsonl. If split is omitted, all available splits are downloaded.",
    )
    output_fpath: Optional[str] = Field(
        default=None,
        description="Exact local file path where the downloaded dataset will be saved. Requires `artifact_fpath` or `split`. Overrides output_dirpath.",
    )
    # Must match `nemo_gym.global_config.HF_TOKEN_KEY_NAME`
    hf_token: Optional[str] = Field(default=None, description="HuggingFace API token for authentication.")
    split: Optional[Literal["train", "validation", "test"]] = Field(
        default=None, description="Dataset split to download. Omit to download all available splits."
    )

    @model_validator(mode="after")
    def check_output_path(self) -> "DownloadJsonlDatasetHuggingFaceConfig":
        if not self.output_dirpath and not self.output_fpath:
            raise ValueError("Either output_dirpath or output_fpath must be provided")
        if self.output_dirpath and self.output_fpath:
            raise ValueError("Cannot specify both output_dirpath and output_fpath")
        if self.artifact_fpath and self.split:
            raise ValueError(
                "Cannot specify both artifact_fpath and split. Use artifact_fpath for targeting a raw file, or split for structured datasets."
            )
        # Prevent output_fpath without split when not using artifact_fpath
        if self.output_fpath and not self.split and not self.artifact_fpath:
            raise ValueError(
                "When using output_fpath without artifact_fpath, split must be specified. Use output_dirpath to download all splits."
            )
        return self


DatasetType = Union[Literal["train"], Literal["validation"], Literal["example"]]


class GitlabDatasetSource(BaseModel):
    """Unified ``source:`` for a dataset fetched from the GitLab model registry."""

    type: Literal["gitlab"]
    dataset_name: str
    version: str
    artifact_fpath: str


class HuggingFaceDatasetSource(BaseModel):
    """Unified ``source:`` for a dataset fetched from the HuggingFace Hub."""

    type: Literal["huggingface"]
    repo_id: str
    artifact_fpath: Optional[str] = None


# One discriminated `source:` block replaces the parallel gitlab_identifier / huggingface_identifier
# fields; `type` selects the backend so it's unambiguous which fields apply.
DatasetSource = Annotated[Union[GitlabDatasetSource, HuggingFaceDatasetSource], Field(discriminator="type")]


class DatasetConfig(BaseModel):
    name: str
    type: DatasetType
    jsonl_fpath: str

    num_repeats: int = Field(default=1, ge=1)
    # Unified, self-describing dataset source. Prefer this over the legacy *_identifier fields below.
    source: Optional[DatasetSource] = None
    # Deprecated: kept working (and back-filled from/into `source`) for backward compatibility.
    gitlab_identifier: Optional[JsonlDatasetGitlabIdentifer] = None
    huggingface_identifier: Optional[JsonlDatasetHuggingFaceIdentifer] = None
    license: Optional[
        Union[
            Literal["Apache 2.0"],
            Literal["MIT"],
            Literal["Creative Commons Attribution 4.0 International"],
            Literal["Creative Commons Attribution-ShareAlike 4.0 International"],
            Literal["NVIDIA Internal Use Only, Do Not Distribute"],
            Literal["NVIDIA Evaluation Dataset License Agreement"],
            Literal["TBD"],
            Literal["GNU General Public License v3.0"],
        ]
    ] = None

    @model_validator(mode="after")
    def check_train_validation_sets(self) -> "DatasetConfig":
        if self.type in ["train", "validation"]:
            assert self.license is not None, f"A license is required for {self.name}"

        return self

    @model_validator(mode="after")
    def normalize_dataset_source(self) -> "DatasetConfig":
        """Reconcile the unified `source:` with the legacy `*_identifier` fields.

        The unified `source:` block is mutually exclusive with the legacy identifiers. The two
        legacy identifiers may still be set together (a gitlab-primary / huggingface-fallback pair
        selected at download time by `config.data_source`) for backward compatibility. A legacy
        identifier emits a deprecation warning and, when a single backend is given, is mirrored into
        `source`; conversely a `source:` is mirrored back into the matching legacy field so existing
        consumers that read `gitlab_identifier`/`huggingface_identifier` keep working.
        """
        legacy_specified = [
            name
            for name, value in (
                ("gitlab_identifier", self.gitlab_identifier),
                ("huggingface_identifier", self.huggingface_identifier),
            )
            if value is not None
        ]
        if self.source is not None and legacy_specified:
            raise ValueError(
                f"Specify a dataset source once for '{self.name}': set only one of "
                f"['source', {', '.join(repr(name) for name in legacy_specified)}]. "
                "Prefer the unified `source:` block."
            )

        if self.source is not None:
            # `source:` was used: back-fill the matching legacy field for existing consumers.
            fields = self.source.model_dump(exclude={"type"})
            if isinstance(self.source, GitlabDatasetSource):
                self.gitlab_identifier = JsonlDatasetGitlabIdentifer(**fields)
            else:
                self.huggingface_identifier = JsonlDatasetHuggingFaceIdentifer(**fields)
            return self

        if not legacy_specified:
            return self

        warnings.warn(
            f"{' and '.join(f'`{name}`' for name in legacy_specified)} "
            f"{'is' if len(legacy_specified) == 1 else 'are'} deprecated for dataset "
            f"'{self.name}'; prefer the unified `source:` block.",
            DeprecationWarning,
            stacklevel=2,
        )
        # Mirror a single legacy identifier into `source`. When both are set (primary + fallback),
        # the single discriminated `source:` can't represent both, so leave it unset and keep the
        # legacy fields as the source of truth.
        if len(legacy_specified) == 1:
            if self.gitlab_identifier is not None:
                self.source = GitlabDatasetSource(type="gitlab", **self.gitlab_identifier.model_dump())
            else:
                self.source = HuggingFaceDatasetSource(type="huggingface", **self.huggingface_identifier.model_dump())

        return self


class BenchmarkDatasetConfig(BaseModel):
    name: str
    type: Literal["benchmark"]
    jsonl_fpath: Path
    prepare_script: Path
    prompt_config: Optional[Path]
    num_repeats: int = Field(default=1, ge=1)


########################################
# Base server config classes
########################################


class Domain(str, Enum):
    """The capability a resources server primarily evaluates or trains.

    Pick the single domain that best fits the task. If several seem to apply, choose the most
    specific one (e.g. prefer `math` or `coding` over `agent`); use `other` only when none
    of the specific values fit. The values:

    - `math`                  — mathematical problem solving (e.g. AIME, MATH, GSM8K).
    - `coding`                — code generation, repair, or execution (e.g. SWE-bench, LiveCodeBench).
    - `agent`                 — multi-step, tool-using / environment-interacting tasks (e.g. tau2,
      workplace_assistant). Prefer a more specific value when the task is really math/coding/etc.
    - `knowledge`             — factual or domain-knowledge question answering (e.g. GPQA, MMLU).
    - `instruction_following` — adherence to explicit formatting/constraints (e.g. IFEval).
    - `long_context`          — reasoning over long inputs (e.g. RULER, long-document QA).
    - `safety`                — refusing harmful content / resisting jailbreaks & prompt injection.
    - `games`                 — interactive game environments (e.g. blackjack, tetris).
    - `translation`           — machine translation quality (e.g. WMT).
    - `e2e`                   — end-to-end pipelines spanning multiple capabilities at once.
    - `rlhf`                  — preference / reward-model / LLM-as-judge evaluations.
    - `other`                 — catch-all when no specific domain above applies.
    """

    MATH = "math"
    CODING = "coding"
    AGENT = "agent"
    KNOWLEDGE = "knowledge"
    INSTRUCTION_FOLLOWING = "instruction_following"
    LONG_CONTEXT = "long_context"
    SAFETY = "safety"
    GAMES = "games"
    TRANSLATION = "translation"
    E2E = "e2e"
    RLHF = "rlhf"
    OTHER = "other"


class BaseServerConfig(BaseModel):
    host: str
    port: int
    num_workers: Optional[int] = None


class BaseRunServerConfig(BaseServerConfig):
    entrypoint: str
    domain: Optional[Domain] = None  # Only required for resources servers


class BaseRunServerInstanceConfig(BaseRunServerConfig):
    name: str  # This name is unique at runtime.


########################################
# Server type and server instance configs
########################################


class BaseRunServerTypeConfig(BaseRunServerConfig):
    model_config = ConfigDict(extra="allow")

    host: Optional[str] = None
    port: Optional[int] = None

    datasets: Optional[List[Union[DatasetConfig, BenchmarkDatasetConfig]]] = None


class BaseServerTypeConfig(BaseModel):
    SERVER_TYPE: ClassVar[
        Union[
            Literal["responses_api_models"],
            Literal["resources_servers"],
            Literal["responses_api_agents"],
        ]
    ]


class ResponsesAPIModelServerTypeConfig(BaseServerTypeConfig):
    SERVER_TYPE: ClassVar[Literal["responses_api_models"]] = "responses_api_models"

    model_config = ConfigDict(extra="allow")

    responses_api_models: Dict[str, BaseRunServerTypeConfig] = Field(min_length=1, max_length=1)


class ResourcesServerTypeConfig(BaseServerTypeConfig):
    SERVER_TYPE: ClassVar[Literal["resources_servers"]] = "resources_servers"

    model_config = ConfigDict(extra="allow")

    resources_servers: Dict[str, BaseRunServerTypeConfig] = Field(min_length=1, max_length=1)


class ResponsesAPIAgentServerTypeConfig(BaseServerTypeConfig):
    SERVER_TYPE: ClassVar[Literal["responses_api_agents"]] = "responses_api_agents"

    model_config = ConfigDict(extra="allow")

    responses_api_agents: Dict[str, BaseRunServerTypeConfig] = Field(min_length=1, max_length=1)


ServerTypeConfig = Union[
    ResponsesAPIModelServerTypeConfig,
    ResourcesServerTypeConfig,
    ResponsesAPIAgentServerTypeConfig,
]


class BaseServerInstanceConfig(BaseServerTypeConfig):
    model_config = ConfigDict(arbitrary_types_allowed=True)

    name: str
    server_type_config_dict: DictConfig = Field(exclude=True)

    @model_validator(mode="after")
    def validate_domain_for_resources_server(self) -> "BaseServerInstanceConfig":
        config = self.get_inner_run_server_config()
        if self.SERVER_TYPE == "resources_servers":
            assert config.domain is not None, "A domain is required for resources servers."
        else:
            # Remove domain field from Model and Agent servers.
            if hasattr(config, "domain"):
                del config.domain
        return self

    def get_server_ref(self) -> ServerRef:
        return is_server_ref({"type": self.SERVER_TYPE, "name": self.name})

    def get_inner_run_server_config_dict(self) -> DictConfig:
        server_type_name = list(getattr(self, self.SERVER_TYPE))[0]
        return self.server_type_config_dict[self.SERVER_TYPE][server_type_name]

    def get_inner_run_server_config(self) -> BaseRunServerTypeConfig:
        return list(getattr(self, self.SERVER_TYPE).values())[0]

    @property
    def datasets(self) -> Optional[List[Union[DatasetConfig, BenchmarkDatasetConfig]]]:
        return self.get_inner_run_server_config().datasets


class ResponsesAPIModelServerInstanceConfig(ResponsesAPIModelServerTypeConfig, BaseServerInstanceConfig):
    pass


class ResourcesServerInstanceConfig(ResourcesServerTypeConfig, BaseServerInstanceConfig):
    pass


class ResponsesAPIAgentServerInstanceConfig(ResponsesAPIAgentServerTypeConfig, BaseServerInstanceConfig):
    pass


ServerInstanceConfig = Union[
    ResponsesAPIModelServerInstanceConfig,
    ResourcesServerInstanceConfig,
    ResponsesAPIAgentServerInstanceConfig,
]
ServerInstanceConfigTypeAdapter = TypeAdapter(ServerInstanceConfig)


def maybe_get_server_instance_config(
    name: str, server_type_config_dict: Any
) -> Tuple[Optional[ServerInstanceConfig], Optional[ValidationError]]:
    """Returns ServerInstanceConfig if a valid server, otherwise None with error details"""
    if not isinstance(server_type_config_dict, DictConfig):
        return None, None

    maybe_server_instance_config_dict = {
        "name": name,
        "server_type_config_dict": server_type_config_dict,
        **OmegaConf.to_container(server_type_config_dict),
    }
    try:
        config = ServerInstanceConfigTypeAdapter.validate_python(maybe_server_instance_config_dict)
        return config, None
    except ValidationError as e:
        return None, e


def is_almost_server(server_type_config_dict: Any) -> bool:
    """Detects if a config looks like a server but might fail validation."""
    from nemo_gym.global_config import ENTRYPOINT_KEY_NAME

    if not isinstance(server_type_config_dict, DictConfig):
        return False

    # Check for server type.
    server_type_keys = ["responses_api_models", "resources_servers", "responses_api_agents"]
    has_server_type = any(key in server_type_config_dict for key in server_type_keys)

    if not has_server_type:
        return False

    # Check for entrypoint presence.
    for server_type_key in server_type_keys:
        if server_type_key in server_type_config_dict:
            inner_dict = server_type_config_dict[server_type_key]
            if isinstance(inner_dict, DictConfig):
                for config in inner_dict.values():
                    if isinstance(config, DictConfig) and ENTRYPOINT_KEY_NAME in config:
                        return True

    return False


########################################
# Train dataset configs
########################################

AGENT_REF_KEY = "agent_ref"


########################################
# Weights and Biases
########################################


class WANDBConfig(BaseModel):
    wandb_project: Optional[str] = None
    wandb_name: Optional[str] = None
    wandb_api_key: Optional[str] = None

    @property
    def is_available(self) -> bool:
        # If global_config recursively hide secrets is called, the api key will be set to ****
        return self.wandb_project and self.wandb_name and self.wandb_api_key and self.wandb_api_key != "****"


########################################
# Weights and Biases
########################################


class AggregateMetricsRequest(BaseModel):
    """POST body for /aggregate_metrics.

    Each item is a stripped verify response dict containing at minimum:
    - TASK_INDEX_KEY_NAME: int
    - "reward": float
    """

    verify_responses: List[Dict[str, Any]]


class AggregateMetrics(BaseModel):
    """Response from /aggregate_metrics.

    Flat string keys for direct logging to W&B/MLflow.
    """

    group_level_metrics: List[Dict[str, Any]] = Field(
        default_factory=list,
        description="Per-task metrics (one dict per task) from RewardProfiler baseline stats.",
    )
    agent_metrics: Dict[str, Any] = Field(
        default_factory=dict,
        description="Overall metrics across all rollouts (RewardProfiler baseline + compute_metrics).",
    )
    key_metrics: Dict[str, Any] = Field(
        default_factory=dict,
        description="Headline metrics for this benchmark. Subset of agent_metrics.",
    )
