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

"""Dynamo runtime configuration ArgGroup."""

from typing import Optional

from dynamo._core import get_reasoning_parser_names, get_tool_parser_names
from dynamo.common.configuration.arg_group import ArgGroup
from dynamo.common.configuration.config_base import ConfigBase
from dynamo.common.configuration.utils import add_argument, add_negatable_bool_argument


class DynamoRuntimeConfig(ConfigBase):
    """Configuration for Dynamo runtime (common across all backends)."""

    namespace: str
    store_kv: str
    request_plane: str
    event_plane: str
    connector: list[str]
    enable_local_indexer: bool
    durable_kv_events: bool

    dyn_tool_call_parser: Optional[str] = None
    dyn_reasoning_parser: Optional[str] = None
    custom_jinja_template: Optional[str] = None
    endpoint_types: str
    dump_config_to: Optional[str] = None

    def validate(self) -> None:
        # TODO  get a better way for spot fixes like this.
        self.enable_local_indexer = not self.durable_kv_events


class DynamoRuntimeArgGroup(ArgGroup):
    """Dynamo runtime configuration parameters (common to all backends)."""

    def add_arguments(self, parser) -> None:
        """Add Dynamo runtime arguments to parser."""
        g = parser.add_argument_group("Dynamo Runtime Options")

        add_argument(
            g,
            flag_name="--namespace",
            env_var="DYN_NAMESPACE",
            default="dynamo",
            help="Dynamo namespace",
        )
        add_argument(
            g,
            flag_name="--store-kv",
            env_var="DYN_STORE_KV",
            default="etcd",
            help="Which key-value backend to use: etcd, mem, file. Etcd uses the ETCD_* env vars (e.g. ETCD_ENDPOINTS) for connection details. File uses root dir from env var DYN_FILE_KV or defaults to $TMPDIR/dynamo_store_kv.",
            choices=["etcd", "file", "mem"],
        )
        add_argument(
            g,
            flag_name="--request-plane",
            env_var="DYN_REQUEST_PLANE",
            default="tcp",
            help="Determines how requests are distributed from routers to workers. 'tcp' is fastest.",
            choices=["tcp", "nats", "http"],
        )
        add_argument(
            g,
            flag_name="--event-plane",
            env_var="DYN_EVENT_PLANE",
            default="nats",
            help="Determines how events are published.",
            choices=["nats", "zmq"],
        )
        add_argument(
            g,
            flag_name="--connector",
            env_var="DYN_CONNECTOR",
            default=["nixl"],
            help="List of connectors to use in order (e.g., --connector nixl lmcache). Options: nixl, lmcache, kvbm, null, none. Order will be preserved in MultiConnector.",
            nargs="*",
        )

        add_negatable_bool_argument(
            g,
            flag_name="--durable-kv-events",
            env_var="DYN_DURABLE_KV_EVENTS",
            default=False,
            help="Enable durable KV events using NATS JetStream instead of the local indexer. By default, local indexer is enabled for lower latency. Use this flag when you need durability and multi-replica router consistency. Requires NATS with JetStream enabled. Can also be set via DYN_DURABLE_KV_EVENTS=true env var.",
        )

        # Optional: tool/reasoning parsers (choices from dynamo._core when available)
        # To avoid name conflicts with different backends, prefix "dyn-" for dynamo specific args
        add_argument(
            g,
            flag_name="--dyn-tool-call-parser",
            env_var="DYN_TOOL_CALL_PARSER",
            default=None,
            help="Tool call parser name for the model.",
            choices=get_tool_parser_names(),
        )
        add_argument(
            g,
            flag_name="--dyn-reasoning-parser",
            env_var="DYN_REASONING_PARSER",
            default=None,
            help="Reasoning parser name for the model. If not specified, no reasoning parsing is performed.",
            choices=get_reasoning_parser_names(),
        )
        add_argument(
            g,
            flag_name="--custom-jinja-template",
            env_var="DYN_CUSTOM_JINJA_TEMPLATE",
            default=None,
            help="Path to a custom Jinja template file to override the model's default chat template. This template will take precedence over any template found in the model repository.",
        )

        add_argument(
            g,
            flag_name="--endpoint-types",
            env_var="DYN_ENDPOINT_TYPES",
            default="chat,completions",
            obsolete_flag="--dyn-endpoint-types",
            help="Comma-separated list of endpoint types to enable. Options: 'chat', 'completions'. Use 'completions' for models without chat templates.",
        )

        add_argument(
            g,
            flag_name="--dump-config-to",
            env_var="DYN_DUMP_CONFIG_TO",
            default=None,
            help="Dump resolved configuration to the specified file path.",
        )
