# SPDX-FileCopyrightText: Copyright (c) 2024-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.
"""
Auto-generated Pydantic models from v1beta1 DGDR Go types.

Generated by: deploy/operator/api/scripts/generate_pydantic_from_go.py
Source: deploy/operator/api/v1beta1/dynamographdeploymentrequest_types.go

DO NOT EDIT MANUALLY - regenerate using the script.
"""

from enum import Enum
from typing import Any, Dict, List, Optional

from pydantic import BaseModel, Field, model_validator

# Import canonical planner types - do NOT redefine them here.
from dynamo.planner.utils.planner_config import (  # noqa: F401 (re-exported)
    PlannerConfig,
    PlannerPreDeploymentSweepMode,
)


class DGDRPhase(str, Enum):
    Pending = "Pending"
    Profiling = "Profiling"
    Ready = "Ready"
    Deploying = "Deploying"
    Deployed = "Deployed"
    Failed = "Failed"


class ProfilingPhase(str, Enum):
    Initializing = "Initializing"
    SweepingPrefill = "SweepingPrefill"
    SweepingDecode = "SweepingDecode"
    SelectingConfig = "SelectingConfig"
    BuildingCurves = "BuildingCurves"
    GeneratingDGD = "GeneratingDGD"
    Done = "Done"


class OptimizationType(str, Enum):
    Latency = "latency"
    Throughput = "throughput"


class SearchStrategy(str, Enum):
    Rapid = "rapid"
    Thorough = "thorough"


class BackendType(str, Enum):
    Auto = "auto"
    Sglang = "sglang"
    Trtllm = "trtllm"
    Vllm = "vllm"


class WorkloadSpec(BaseModel):
    """WorkloadSpec defines the workload characteristics for SLA-based profiling."""

    isl: Optional[int] = Field(
        default=4000, description="ISL is the Input Sequence Length (number of tokens)."
    )
    osl: Optional[int] = Field(
        default=1000,
        description="OSL is the Output Sequence Length (number of tokens).",
    )
    concurrency: Optional[float] = Field(
        default=None,
        description="Concurrency is the target concurrency level. Required (or RequestRate) when the planner is disabled.",
    )
    requestRate: Optional[float] = Field(
        default=None,
        description="RequestRate is the target request rate (req/s). Required (or Concurrency) when the planner is disabled.",
    )


class SLASpec(BaseModel):
    """Service-level agreement targets.

    Provide exactly one of:

    - ``ttft`` + ``itl``: explicit latency targets (default: 2000 ms / 30 ms)
    - ``e2eLatency``: end-to-end latency target
    - ``optimizationType``: high-level objective without explicit numeric targets"""

    optimizationType: Optional[OptimizationType] = Field(
        default=None,
        description="OptimizationType controls the profiling optimization strategy. Use when explicit SLA targets (ttft+itl or e2eLatency) are not known.",
    )
    ttft: Optional[float] = Field(
        default=2000,
        description="TTFT is the Time To First Token target in milliseconds.",
    )
    itl: Optional[float] = Field(
        default=30, description="ITL is the Inter-Token Latency target in milliseconds."
    )
    e2eLatency: Optional[float] = Field(
        default=None,
        description="E2ELatency is the target end-to-end request latency in milliseconds. Alternative to specifying TTFT + ITL.",
    )

    @model_validator(mode="after")
    def _validate_sla_options(self) -> "SLASpec":
        """Ensure at most one SLA mode is active."""
        has_e2e = self.e2eLatency is not None
        has_opt = self.optimizationType is not None
        ttft_itl_touched = (
            "ttft" in self.model_fields_set or "itl" in self.model_fields_set
        )
        has_ttft_itl = (self.ttft is not None and self.itl is not None) and (
            ttft_itl_touched or (not has_e2e and not has_opt)
        )
        options_count = sum([has_ttft_itl, has_e2e, has_opt])
        if options_count > 1:
            raise ValueError(
                "SLA must specify exactly one of: (ttft and itl), e2eLatency, "
                "or optimizationType — not multiple."
            )
        if (self.ttft is not None) != (self.itl is not None):
            raise ValueError("ttft and itl must both be provided together.")
        return self


class ModelCacheSpec(BaseModel):
    """ModelCacheSpec references a PVC containing pre-downloaded model weights."""

    pvcName: Optional[str] = Field(
        default=None,
        description="PVCName is the name of the PersistentVolumeClaim containing model weights. The PVC must exist in the same namespace as the DGDR.",
    )
    pvcModelPath: Optional[str] = Field(
        default=None,
        description='PVCModelPath is the path to the model checkpoint directory within the PVC (e.g. "deepseek-r1" or "models/Llama-3.1-405B-FP8").',
    )
    pvcMountPath: str = Field(
        default="/opt/model-cache",
        description="PVCMountPath is the mount path for the PVC inside the container.",
    )


class OverridesSpec(BaseModel):
    """OverridesSpec allows customizing the profiling job and the generated DynamoGraphDeployment."""

    profilingJob: Optional[Dict[str, Any]] = Field(
        default=None,
        description="ProfilingJob allows overriding the profiling Job specification. Fields set here are merged into the controller-generated Job spec.",
    )
    dgd: Optional[Dict[str, Any]] = Field(
        default=None,
        description="DGD allows providing a full or partial nvidia.com/v1alpha1 DynamoGraphDeployment to use as the base for the generated deployment. Fields from profiling results are merged on top. Use this to override backend worker images.  The field is stored as a raw embedded resource rather than a typed *v1alpha1.DynamoGraphDeployment to avoid a circular import: v1alpha1 already imports v1beta1 as the conversion hub and Go does not allow import cycles.  The EmbeddedResource marker tells the API server to validate that the value is a well-formed Kubernetes object (has apiVersion/kind), but does not enforce that it is specifically a DynamoGraphDeployment. Full type validation (correct apiVersion, kind, and field schema) is performed by the controller during reconciliation. TODO(future MR): add webhook admission validation for the DGD field type.",
    )


class MockerSpec(BaseModel):
    """MockerSpec configures the simulated (mocker) backend."""

    enabled: Optional[bool] = Field(
        default=None,
        description="Enabled indicates whether to deploy mocker workers instead of real inference workers. Useful for large-scale testing without GPUs.",
    )


class KVRouterSpec(BaseModel):
    """KVRouterSpec configures KV-cache-aware routing."""

    enabled: Optional[bool] = Field(
        default=None,
        description="Enabled indicates whether to enable KV-cache-aware routing in the generated DGD. KV routing optimizes request scheduling based on KV cache locality.",
    )


class FeaturesSpec(BaseModel):
    """FeaturesSpec controls optional Dynamo platform features in the generated deployment."""

    planner: Optional[PlannerConfig] = Field(
        default=None,
        description="Planner is the raw SLA planner configuration passed to the planner service. Its schema is defined by dynamo.planner.utils.planner_config.PlannerConfig. Go treats this as opaque bytes; the Planner service validates it at startup. The presence of this field (non-null) enables the planner in the generated DGD.",
    )
    mocker: Optional[MockerSpec] = Field(
        default=None,
        description="Mocker configures the simulated (mocker) backend for testing without GPUs.",
    )


class HardwareSpec(BaseModel):
    """HardwareSpec describes the hardware resources available for profiling and deployment. These fields are typically auto-filled by the operator from cluster discovery."""

    gpuSku: Optional[str] = Field(
        default=None,
        description='GPUSKU is the GPU SKU identifier (e.g., "H100_SXM", "A100_80GB").',
    )
    vramMb: Optional[float] = Field(
        default=None, description="VRAMMB is the VRAM per GPU in MiB."
    )
    totalGpus: Optional[int] = Field(
        default=None,
        description="TotalGPUs is the total number of GPUs available in the cluster.",
    )
    numGpusPerNode: Optional[int] = Field(
        default=None, description="NumGPUsPerNode is the number of GPUs per node."
    )


class DynamoGraphDeploymentRequestSpec(BaseModel):
    """DynamoGraphDeploymentRequestSpec defines the desired state of a DynamoGraphDeploymentRequest. Only the Model field is required; all other fields are optional and have sensible defaults."""

    model: str = Field(
        description='Model specifies the model to deploy (e.g., "Qwen/Qwen3-0.6B", "meta-llama/Llama-3-70b"). Can be a HuggingFace ID or a private model name.'
    )
    backend: BackendType = Field(
        default="auto",
        description="Backend specifies the inference backend to use for profiling and deployment.",
    )
    image: Optional[str] = Field(
        default=None,
        description='Image is the container image reference for the profiling job (frontend image). Example: "nvcr.io/nvidia/ai-dynamo/dynamo-frontend:1.0.0".',
    )
    modelCache: Optional[ModelCacheSpec] = Field(
        default=None,
        description="ModelCache provides optional PVC configuration for pre-downloaded model weights. When provided, weights are loaded from the PVC instead of downloading from HuggingFace.",
    )
    hardware: Optional[HardwareSpec] = Field(
        default=None,
        description="Hardware describes the hardware resources available for profiling and deployment. Typically auto-filled by the operator from cluster discovery.",
    )
    workload: Optional[WorkloadSpec] = Field(
        default=None,
        description="Workload defines the expected workload characteristics for SLA-based profiling.",
    )
    sla: Optional[SLASpec] = Field(
        default=None,
        description="SLA defines service-level agreement targets that drive profiling optimization.",
    )
    overrides: Optional[OverridesSpec] = Field(
        default=None,
        description="Overrides allows customizing the profiling job and the generated DynamoGraphDeployment.",
    )
    features: Optional[FeaturesSpec] = Field(
        default=None,
        description="Features controls optional Dynamo platform features in the generated deployment.",
    )
    searchStrategy: SearchStrategy = Field(
        default="rapid",
        description='SearchStrategy controls the profiling search depth. "rapid" performs a fast sweep; "thorough" explores more configurations.',
    )
    autoApply: Optional[bool] = Field(
        default=True,
        description="AutoApply indicates whether to automatically create a DynamoGraphDeployment after profiling completes. If false, the generated spec is stored in status for manual review and application.",
    )


class ParetoConfig(BaseModel):
    """ParetoConfig represents a single Pareto-optimal deployment configuration discovered during profiling."""

    config: Dict[str, Any] = Field(
        description="Config is the full deployment configuration for this Pareto point."
    )


class ProfilingResultsStatus(BaseModel):
    """ProfilingResultsStatus contains the output of the profiling process."""

    pareto: Optional[List[ParetoConfig]] = Field(
        default=None,
        description="Pareto is the list of Pareto-optimal deployment configurations discovered during profiling. Each entry represents a different cost/performance trade-off.",
    )
    selectedConfig: Optional[Dict[str, Any]] = Field(
        default=None,
        description="SelectedConfig is the recommended configuration chosen by the profiler based on the SLA targets. This is the configuration used for deployment when autoApply is true.",
    )


class DeploymentInfoStatus(BaseModel):
    """DeploymentInfoStatus tracks the state of the deployed DynamoGraphDeployment."""

    replicas: Optional[int] = Field(
        default=None, description="Replicas is the desired number of replicas."
    )
    availableReplicas: Optional[int] = Field(
        default=None,
        description="AvailableReplicas is the number of replicas that are available and ready.",
    )


class DynamoGraphDeploymentRequestStatus(BaseModel):
    """DynamoGraphDeploymentRequestStatus represents the observed state of a DynamoGraphDeploymentRequest."""

    phase: Optional[DGDRPhase] = Field(
        default=None,
        description="Phase is the high-level lifecycle phase of the deployment request.",
    )
    profilingPhase: Optional[ProfilingPhase] = Field(
        default=None,
        description='ProfilingPhase indicates the current sub-phase of the profiling pipeline. Only meaningful when Phase is "Profiling". Cleared when profiling completes or fails.',
    )
    dgdName: Optional[str] = Field(
        default=None,
        description="DGDName is the name of the generated or created DynamoGraphDeployment.",
    )
    profilingJobName: Optional[str] = Field(
        default=None,
        description="ProfilingJobName is the name of the Kubernetes Job running the profiler.",
    )
    profilingResults: Optional[ProfilingResultsStatus] = Field(
        default=None,
        description="ProfilingResults contains the output of the profiling process including Pareto-optimal configurations and the selected deployment configuration.",
    )
    deploymentInfo: Optional[DeploymentInfoStatus] = Field(
        default=None,
        description="DeploymentInfo tracks the state of the deployed DynamoGraphDeployment. Populated when a DGD has been created (either via autoApply or manually).",
    )
    observedGeneration: Optional[int] = Field(
        default=None,
        description="ObservedGeneration is the most recent generation observed by the controller.",
    )


class DynamoGraphDeploymentRequest(BaseModel):
    """DynamoGraphDeploymentRequest is the Schema for the dynamographdeploymentrequests API. It provides a simplified, SLA-driven interface for deploying inference models on Dynamo. Users specify a model and optional performance targets; the controller handles profiling, configuration selection, and deployment."""

    spec: Optional[DynamoGraphDeploymentRequestSpec] = Field(
        default=None,
        description="Spec defines the desired state for this deployment request.",
    )
    status: Optional[DynamoGraphDeploymentRequestStatus] = Field(
        default=None,
        description="Status reflects the current observed state of this deployment request.",
    )
