# SPDX-FileCopyrightText: Copyright (c) 2025-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.

"""Shared helpers and configuration for the profiler pipeline."""

import copy
import logging
import os
from dataclasses import dataclass, field

import pandas as pd

from dynamo.planner.config.planner_config import PlannerPreDeploymentSweepMode
from dynamo.profiler.utils.config_modifiers.parallelization_mapping import (
    PickedParallelConfig,
)
from dynamo.profiler.utils.dgdr_v1beta1_types import (
    DynamoGraphDeploymentRequestSpec,
    ProfilingPhase,
)

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Published container image naming conventions
# ---------------------------------------------------------------------------

# Mapping from backend name to the image-name component of the published
# backend runtime image.
# e.g. vllm → nvcr.io/nvidia/ai-dynamo/vllm-runtime:1.0.0
BACKEND_IMAGE_NAMES: dict[str, str] = {
    "vllm": "vllm-runtime",
    "sglang": "sglang-runtime",
    "trtllm": "tensorrtllm-runtime",
}

PLANNER_IMAGE_NAME = "dynamo-planner"


def _replace_image_name(image_ref: str, new_name: str) -> str:
    """Replace the image name component in a Docker image reference.

    Preserves the registry path prefix and tag suffix, only replacing the
    last ``/``-delimited component (before any ``:tag``).
    """
    slash_idx = image_ref.rfind("/")
    prefix = image_ref[: slash_idx + 1] if slash_idx >= 0 else ""
    suffix = image_ref[slash_idx + 1 :]
    name_and_tag, has_digest, digest = suffix.partition("@")
    colon_idx = name_and_tag.rfind(":")
    tag = name_and_tag[colon_idx:] if colon_idx >= 0 else ""
    digest_suffix = f"@{digest}" if has_digest else ""
    return f"{prefix}{new_name}{tag}{digest_suffix}"


def derive_planner_image(profiler_image: str) -> str:
    """Derive the planner service image from the profiler image reference."""
    return _replace_image_name(profiler_image, PLANNER_IMAGE_NAME)


def derive_backend_image(profiler_image: str, backend: str) -> str:
    """Derive the backend worker image from the profiler image.

    Replaces the image name (the last ``/``-delimited component, before any
    ``:tag``) with the backend-specific runtime image name, preserving the
    registry path and tag unchanged.

    Examples::

        derive_backend_image(
            "nvcr.io/nvidia/ai-dynamo/dynamo-frontend:1.0.0", "vllm"
        )
        # → "nvcr.io/nvidia/ai-dynamo/vllm-runtime:1.0.0"

        derive_backend_image("myregistry.io/sglang-runtime:1.0.0", "sglang")
        # → "myregistry.io/sglang-runtime:1.0.0"

    Args:
        profiler_image: Any Docker image reference of the form
            ``[REGISTRY/]NAME[:TAG]``.
        backend: The resolved backend type (``'vllm'``, ``'sglang'``, or
            ``'trtllm'``).

    Returns:
        The backend container image string.

    Raises:
        ValueError: If *backend* is not a recognised backend.
    """
    backend_image_name = BACKEND_IMAGE_NAMES.get(backend)
    if backend_image_name is None:
        raise ValueError(
            f"Cannot derive backend image for unknown backend '{backend}'. "
            f"Supported backends: {list(BACKEND_IMAGE_NAMES.keys())}"
        )

    return _replace_image_name(profiler_image, backend_image_name)


# ---------------------------------------------------------------------------
# Operational defaults not part of DynamoGraphDeploymentRequestSpec
# ---------------------------------------------------------------------------

DEFAULT_OUTPUT_DIR = "profiling_results"
DEFAULT_NAMESPACE = os.environ.get("DGDR_NAMESPACE", "dynamo-sla-profiler")
DEFAULT_DEPLOYMENT_TIMEOUT = 3600
DEFAULT_PREFILL_INTERPOLATION_GRANULARITY = 16
DEFAULT_DECODE_INTERPOLATION_GRANULARITY = 6
DEFAULT_DRY_RUN = False


@dataclass
class ProfilerOperationalConfig:
    """Operational knobs that are not part of the DGDR spec."""

    output_dir: str = DEFAULT_OUTPUT_DIR
    k8s_namespace: str = DEFAULT_NAMESPACE
    deployment_timeout: int = DEFAULT_DEPLOYMENT_TIMEOUT
    prefill_interpolation_granularity: int = DEFAULT_PREFILL_INTERPOLATION_GRANULARITY
    decode_interpolation_granularity: int = DEFAULT_DECODE_INTERPOLATION_GRANULARITY
    dry_run: bool = DEFAULT_DRY_RUN
    current_phase: ProfilingPhase = field(default=ProfilingPhase.Initializing)


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def picked_config_from_row(prefix: str, row: pd.Series) -> PickedParallelConfig:
    """Extract a PickedParallelConfig from a picked ColumnsDisagg DataFrame row."""
    return PickedParallelConfig(
        tp=int(row.get(f"{prefix}tp", 1)),
        pp=int(row.get(f"{prefix}pp", 1)),
        dp=int(row.get(f"{prefix}dp", 1)),
        moe_tp=int(row.get(f"{prefix}moe_tp", 1)),
        moe_ep=int(row.get(f"{prefix}moe_ep", 1)),
    )


def resolve_model_path(dgdr: DynamoGraphDeploymentRequestSpec) -> str:
    """Resolve the model path, preferring local PVC mount over HF ID."""
    if (
        dgdr.modelCache
        and dgdr.modelCache.pvcName
        and dgdr.modelCache.pvcMountPath
        and dgdr.modelCache.pvcModelPath
    ):
        mount = dgdr.modelCache.pvcMountPath.rstrip("/")
        sub = dgdr.modelCache.pvcModelPath.strip("/")
        local_path = f"{mount}/{sub}"
        if os.path.isdir(local_path):
            return local_path
    return dgdr.model


def is_planner_enabled(dgdr: DynamoGraphDeploymentRequestSpec) -> bool:
    """True when the DGDR spec has a planner config with scaling enabled."""
    return (
        dgdr.features is not None
        and dgdr.features.planner is not None
        and dgdr.features.planner.scaling_enabled()
    )


def is_mocker_enabled(dgdr: DynamoGraphDeploymentRequestSpec) -> bool:
    """True when the DGDR spec has mocker explicitly enabled."""
    return (
        dgdr.features is not None
        and dgdr.features.mocker is not None
        and dgdr.features.mocker.enabled is True
    )


def needs_profile_data(dgdr: DynamoGraphDeploymentRequestSpec) -> bool:
    """True when the DGDR requires profiling interpolation data *at this stage*.

    Profile data (NPZ/JSON on disk) is consumed by:

    * **Mocker workers** for latency simulation — required for thorough
      mode. In rapid mode the mocker pulls latency data directly from the
      AIConfigurator SDK via ``--aic-perf-model`` flags injected by the
      profiler, so no NPZ is emitted.
    * **Planner** when throughput scaling is enabled — required for
      thorough mode only. In rapid mode the planner runs AIC interpolation
      in-process at bootstrap (see ``aic_interpolation.py``), so the
      profiler no longer emits NPZ for planner rapid deployments either.
    """
    sweep_mode = (
        dgdr.features.planner.pre_deployment_sweeping_mode
        if dgdr.features is not None and dgdr.features.planner is not None
        else None
    )
    is_rapid = sweep_mode == PlannerPreDeploymentSweepMode.Rapid
    if is_mocker_enabled(dgdr):
        return not is_rapid
    if (
        dgdr.features is not None
        and dgdr.features.planner is not None
        and dgdr.features.planner.enable_throughput_scaling
    ):
        return not is_rapid
    return False


def determine_picking_mode(dgdr: DynamoGraphDeploymentRequestSpec) -> str:
    target_load_provided = dgdr.workload is not None and (
        dgdr.workload.requestRate is not None or dgdr.workload.concurrency is not None
    )
    if is_planner_enabled(dgdr):
        return "autoscale"
    elif target_load_provided:
        return "load_match"
    return "default"


def warn_and_update_sla(
    best_latencies: dict,
    target_ttft: float,
    target_tpot: float,
) -> tuple[float, float]:
    """Warn if SLA is unachievable; return (possibly updated) targets."""
    achieved_ttft = best_latencies.get("ttft", 0.0)
    achieved_tpot = best_latencies.get("tpot", 0.0)

    if achieved_ttft > target_ttft:
        logger.warning(
            "TTFT SLA %.1fms is unachievable. Best achievable: %.1fms. Updating SLA.",
            target_ttft,
            achieved_ttft,
        )
        target_ttft = achieved_ttft

    if achieved_tpot > target_tpot:
        logger.warning(
            "ITL SLA %.1fms is unachievable. Best achievable: %.1fms. Updating SLA.",
            target_tpot,
            achieved_tpot,
        )
        target_tpot = achieved_tpot

    return target_ttft, target_tpot


def warn_gpu_shortage(
    picking_mode: str,
    best_latencies: dict,
    total_gpus: int,
) -> None:
    if picking_mode != "load_match":
        return
    gpus_needed = best_latencies.get("total_gpus_needed")
    if gpus_needed is not None and gpus_needed > total_gpus:
        logger.warning(
            "Load target requires %d GPUs but only %d available. "
            "Consider adding more GPUs or reducing the load target.",
            gpus_needed,
            total_gpus,
        )


def get_profiling_job_tolerations(dgdr: DynamoGraphDeploymentRequestSpec) -> list:
    """Return tolerations from overrides.profilingJob.template.spec.tolerations."""
    try:
        if dgdr.overrides is None or dgdr.overrides.profilingJob is None:
            return []
        return (
            dgdr.overrides.profilingJob.get("template", {})
            .get("spec", {})
            .get("tolerations", [])
        )
    except (AttributeError, KeyError):
        return []


def inject_tolerations_into_dgd(dgd_config: dict, tolerations: list) -> dict:
    """Add tolerations to every service's extraPodSpec in a DGD config dict.

    Tolerations already present in a service are preserved; only new entries
    (by identity) are appended.  Returns a deep copy with tolerations applied.
    """
    result = copy.deepcopy(dgd_config)
    for _svc_name, svc in result.get("spec", {}).get("services", {}).items():
        if not isinstance(svc, dict):
            continue
        eps = svc.setdefault("extraPodSpec", {})
        existing = eps.get("tolerations", [])
        new_entries = [t for t in tolerations if t not in existing]
        if new_entries:
            eps["tolerations"] = list(existing) + new_entries
    return result
