# 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.

import os
from unittest.mock import AsyncMock, Mock, call, patch

import pytest

from dynamo.planner.config.defaults import SubComponentType, TargetReplica
from dynamo.planner.connectors.kubernetes import KubernetesConnector
from dynamo.planner.errors import (
    DeploymentModelNameMismatchError,
    DeploymentValidationError,
    DuplicateSubComponentError,
    DynamoGraphDeploymentNotFoundError,
    EmptyTargetReplicasError,
    ModelNameNotFoundError,
    SubComponentNotFoundError,
)
from dynamo.planner.monitoring.dgd_services import (
    Service,
    get_service_from_sub_component_type_or_name,
)

pytestmark = [
    pytest.mark.gpu_0,
    pytest.mark.pre_merge,
    pytest.mark.unit,
    pytest.mark.planner,
]


@pytest.fixture
def mock_kube_api():
    mock_api = Mock()
    mock_api.get_graph_deployment = Mock()
    mock_api.update_graph_replicas = AsyncMock()
    mock_api.wait_for_graph_deployment_ready = AsyncMock()
    mock_api.is_deployment_ready = Mock()
    return mock_api


@pytest.fixture
def mock_kube_api_class(mock_kube_api):
    mock_class = Mock()
    mock_class.return_value = mock_kube_api
    return mock_class


@pytest.fixture
def kubernetes_connector(mock_kube_api_class, monkeypatch):
    # Patch the KubernetesAPI class before instantiating the connector
    monkeypatch.setattr(
        "dynamo.planner.connectors.kubernetes.KubernetesAPI", mock_kube_api_class
    )
    with patch.dict(os.environ, {"DYN_PARENT_DGD_K8S_NAME": "test-graph"}):
        connector = KubernetesConnector("test-dynamo-namespace")
        return connector


def test_kubernetes_connector_no_env_var():
    with patch("dynamo.planner.connectors.kubernetes.KubernetesAPI"):
        with pytest.raises(DeploymentValidationError) as exc_info:
            KubernetesConnector("test-dynamo-namespace")

    exception = exc_info.value
    assert set(exception.errors) == {
        "DYN_PARENT_DGD_K8S_NAME environment variable is not set"
    }


def test_get_service_name_from_sub_component_type(kubernetes_connector):
    deployment = {
        "metadata": {"name": "test-graph"},
        "spec": {
            "services": {
                "test-component-prefill": {
                    "replicas": 2,
                    "subComponentType": "prefill",
                },
                "test-component-decode": {"replicas": 3, "subComponentType": "decode"},
            }
        },
    }

    service = get_service_from_sub_component_type_or_name(
        deployment, SubComponentType.PREFILL
    )
    assert service.name == "test-component-prefill"
    assert service.number_replicas() == 2

    # should still work if the component_name is provided
    service = get_service_from_sub_component_type_or_name(
        deployment, SubComponentType.PREFILL, "test-component-prefill"
    )
    assert service.name == "test-component-prefill"
    assert service.number_replicas() == 2

    # should respect subComponentType first
    service = get_service_from_sub_component_type_or_name(
        deployment, SubComponentType.DECODE, "test-component-prefill"
    )
    assert service.name == "test-component-decode"
    assert service.number_replicas() == 3


def test_get_service_name_from_sub_component_type_not_found(kubernetes_connector):
    deployment = {
        "metadata": {"name": "test-graph"},
        "spec": {
            "services": {
                "test-component-decode": {"replicas": 3, "subComponentType": "decode"},
            }
        },
    }
    with pytest.raises(SubComponentNotFoundError) as exc_info:
        get_service_from_sub_component_type_or_name(
            deployment, SubComponentType.PREFILL
        )

    with pytest.raises(SubComponentNotFoundError) as exc_info:
        get_service_from_sub_component_type_or_name(
            deployment, SubComponentType.PREFILL, "test-component-decode"
        )

    exception = exc_info.value
    assert exception.sub_component_type == SubComponentType.PREFILL.value


def test_get_service_name_from_sub_component_type_duplicate(kubernetes_connector):
    deployment = {
        "metadata": {"name": "test-graph"},
        "spec": {
            "services": {
                "test-component-prefill": {
                    "replicas": 2,
                    "subComponentType": "prefill",
                },
                "test-component-prefill-2": {
                    "replicas": 3,
                    "subComponentType": "prefill",
                },
            }
        },
    }

    with pytest.raises(DuplicateSubComponentError) as exc_info:
        # even though "test-component-prefill" is provided, subComponentType duplicates should result in an error
        get_service_from_sub_component_type_or_name(
            deployment, SubComponentType.PREFILL, "test-component-prefill"
        )

    exception = exc_info.value
    assert exception.sub_component_type == SubComponentType.PREFILL.value
    assert set(exception.service_names) == {
        "test-component-prefill",
        "test-component-prefill-2",
    }


def test_get_service_name_from_sub_component_type_or_name(kubernetes_connector):
    deployment = {
        "metadata": {"name": "test-graph"},
        "spec": {
            "services": {
                "test-component-prefill": {"replicas": 2},
                "test-component-decode": {"replicas": 3},
            }
        },
    }

    service = get_service_from_sub_component_type_or_name(
        deployment, SubComponentType.PREFILL, "test-component-prefill"
    )
    assert service.name == "test-component-prefill"
    assert service.number_replicas() == 2


@pytest.mark.asyncio
async def test_add_component_increases_replicas(kubernetes_connector, mock_kube_api):
    # Arrange
    sub_component_type = SubComponentType.PREFILL
    component_name = "test-component"
    mock_deployment = {
        "metadata": {"name": "test-graph"},
        "spec": {
            "services": {
                component_name: {
                    "replicas": 1,
                    "subComponentType": sub_component_type.value,
                }
            }
        },
    }
    mock_kube_api.get_graph_deployment.return_value = mock_deployment
    mock_kube_api.update_graph_replicas.return_value = None
    mock_kube_api.wait_for_graph_deployment_ready.return_value = None

    # Act
    await kubernetes_connector.add_component(sub_component_type)

    # Assert
    mock_kube_api.get_graph_deployment.assert_called_once()
    mock_kube_api.update_graph_replicas.assert_called_once_with(
        "test-graph", component_name, 2
    )
    mock_kube_api.wait_for_graph_deployment_ready.assert_called_once_with("test-graph")


@pytest.mark.asyncio
async def test_add_component_with_no_replicas_specified(
    kubernetes_connector, mock_kube_api
):
    # Arrange
    sub_component_type = SubComponentType.PREFILL
    component_name = "test-component"
    mock_deployment = {
        "metadata": {"name": "test-graph"},
        "spec": {
            "services": {component_name: {"subComponentType": sub_component_type.value}}
        },
    }
    mock_kube_api.get_graph_deployment.return_value = mock_deployment

    # Act
    await kubernetes_connector.add_component(sub_component_type)

    # Assert
    mock_kube_api.update_graph_replicas.assert_called_once_with(
        "test-graph", component_name, 1
    )
    mock_kube_api.wait_for_graph_deployment_ready.assert_called_once_with("test-graph")


@pytest.mark.asyncio
async def test_add_component_deployment_not_found(kubernetes_connector, mock_kube_api):
    # Arrange
    component_name = "test-component"
    mock_kube_api.get_graph_deployment.side_effect = DynamoGraphDeploymentNotFoundError(
        "test-graph", "default"
    )

    # Act & Assert
    with pytest.raises(DynamoGraphDeploymentNotFoundError):
        await kubernetes_connector.add_component(component_name)


@pytest.mark.asyncio
async def test_add_component_component_not_found(kubernetes_connector, mock_kube_api):
    # Arrange
    mock_deployment = {
        "metadata": {"name": "test-graph"},
        "spec": {"services": {"test-component": {"subComponentType": "decode"}}},
    }
    mock_kube_api.get_graph_deployment.return_value = mock_deployment

    # Act
    with pytest.raises(SubComponentNotFoundError) as exc_info:
        await kubernetes_connector.add_component(SubComponentType.PREFILL)

        mock_kube_api.update_graph_replicas.assert_not_called()
        mock_kube_api.wait_for_graph_deployment_ready.assert_not_called()

    exception = exc_info.value
    assert exception.sub_component_type == "prefill"


@pytest.mark.asyncio
async def test_remove_component_decreases_replicas(kubernetes_connector, mock_kube_api):
    # Arrange
    component_name = "test-component"
    sub_component_type = SubComponentType.PREFILL
    mock_deployment = {
        "metadata": {"name": "test-graph"},
        "spec": {
            "services": {
                "test-component": {
                    "replicas": 2,
                    "subComponentType": sub_component_type.value,
                }
            }
        },
    }
    mock_kube_api.get_graph_deployment.return_value = mock_deployment

    # Act
    await kubernetes_connector.remove_component(sub_component_type)

    # Assert
    mock_kube_api.update_graph_replicas.assert_called_once_with(
        "test-graph", component_name, 1
    )
    mock_kube_api.wait_for_graph_deployment_ready.assert_called_once_with("test-graph")


@pytest.mark.asyncio
async def test_remove_component_with_zero_replicas(kubernetes_connector, mock_kube_api):
    # Arrange
    component_name = "test-component"
    sub_component_type = SubComponentType.PREFILL
    mock_deployment = {
        "metadata": {"name": "test-graph"},
        "spec": {
            "services": {
                component_name: {
                    "replicas": 0,
                    "subComponentType": sub_component_type.value,
                }
            }
        },
    }
    mock_kube_api.get_graph_deployment.return_value = mock_deployment

    # Act
    await kubernetes_connector.remove_component(sub_component_type)

    # Assert
    mock_kube_api.update_graph_replicas.assert_not_called()
    mock_kube_api.wait_for_graph_deployment_ready.assert_not_called()


@pytest.mark.asyncio
async def test_remove_component_component_not_found(
    kubernetes_connector, mock_kube_api
):
    # Arrange
    component_name = "test-component"
    sub_component_type = SubComponentType.PREFILL
    mock_deployment = {
        "metadata": {"name": "test-graph"},
        "spec": {
            "services": {
                component_name: {
                    "replicas": 0,
                    "subComponentType": sub_component_type.value,
                }
            }
        },
    }
    mock_kube_api.get_graph_deployment.return_value = mock_deployment

    # Act
    with pytest.raises(SubComponentNotFoundError) as exc_info:
        await kubernetes_connector.remove_component(SubComponentType.DECODE)

        # Assert
        mock_kube_api.update_graph_replicas.assert_not_called()
        mock_kube_api.wait_for_graph_deployment_ready.assert_not_called()

    exception = exc_info.value
    assert exception.sub_component_type == "decode"


@pytest.mark.asyncio
async def test_set_component_replicas(kubernetes_connector, mock_kube_api):
    # Arrange
    target_replicas = [
        TargetReplica(sub_component_type=SubComponentType.PREFILL, desired_replicas=3),
        TargetReplica(
            sub_component_type=SubComponentType.DECODE,
            component_name="component2",
            desired_replicas=2,
        ),
    ]
    mock_deployment = {
        "metadata": {"name": "test-graph"},
        "spec": {
            "services": {
                "component1": {"replicas": 1, "subComponentType": "prefill"},
                "component2": {"replicas": 1},
            }
        },
    }
    mock_kube_api.get_graph_deployment.return_value = mock_deployment
    mock_kube_api.is_deployment_ready.return_value = True
    mock_kube_api.wait_for_graph_deployment_ready.return_value = None

    # Act
    await kubernetes_connector.set_component_replicas(target_replicas)

    # Assert
    mock_kube_api.get_graph_deployment.assert_called_once()
    mock_kube_api.is_deployment_ready.assert_called_once_with(mock_deployment)
    # Should be called twice, once for each component
    expected_calls = [
        call("test-graph", "component1", 3),  # prefill component with 3 replicas
        call("test-graph", "component2", 2),  # decode component with 2 replicas
    ]
    mock_kube_api.update_graph_replicas.assert_has_calls(expected_calls, any_order=True)
    mock_kube_api.wait_for_graph_deployment_ready.assert_called_once_with("test-graph")


@pytest.mark.asyncio
async def test_set_component_replicas_component_not_found(
    kubernetes_connector, mock_kube_api
):
    # Arrange
    target_replicas = [
        TargetReplica(sub_component_type=SubComponentType.PREFILL, desired_replicas=3),
        TargetReplica(sub_component_type=SubComponentType.DECODE, desired_replicas=2),
    ]
    mock_deployment = {
        "metadata": {"name": "test-graph"},
        "spec": {
            "services": {
                "component1": {"replicas": 1, "subComponentType": "prefill"},
                "component2": {"replicas": 1},
            }
        },
    }
    mock_kube_api.get_graph_deployment.return_value = mock_deployment
    mock_kube_api.is_deployment_ready.return_value = True
    mock_kube_api.update_graph_replicas.return_value = None
    mock_kube_api.wait_for_graph_deployment_ready.return_value = None

    # Act
    with pytest.raises(SubComponentNotFoundError) as exc_info:
        await kubernetes_connector.set_component_replicas(target_replicas)

    exception = exc_info.value
    assert exception.sub_component_type == SubComponentType.DECODE.value


@pytest.mark.asyncio
async def test_set_component_replicas_component_already_at_desired_replicas(
    kubernetes_connector, mock_kube_api
):
    # Arrange
    target_replicas = [
        TargetReplica(sub_component_type=SubComponentType.PREFILL, desired_replicas=3),
        TargetReplica(sub_component_type=SubComponentType.DECODE, desired_replicas=2),
    ]
    mock_deployment = {
        "metadata": {"name": "test-graph"},
        "spec": {
            "services": {
                "component1": {"replicas": 1, "subComponentType": "prefill"},
                "component2": {"replicas": 2, "subComponentType": "decode"},
            }
        },
    }
    mock_kube_api.get_graph_deployment.return_value = mock_deployment
    mock_kube_api.is_deployment_ready.return_value = True
    mock_kube_api.update_graph_replicas.return_value = None
    mock_kube_api.wait_for_graph_deployment_ready.return_value = None

    # Act
    await kubernetes_connector.set_component_replicas(target_replicas)

    # Assert
    mock_kube_api.get_graph_deployment.assert_called_once()
    mock_kube_api.is_deployment_ready.assert_called_once_with(mock_deployment)

    # Should be called once, for the prefill component (decode component is already at desired replicas)
    mock_kube_api.update_graph_replicas.assert_called_once_with(
        "test-graph", "component1", 3
    )
    mock_kube_api.wait_for_graph_deployment_ready.assert_called_once_with("test-graph")


@pytest.mark.asyncio
async def test_set_component_replicas_deployment_not_found(
    kubernetes_connector, mock_kube_api
):
    # Arrange
    target_replicas = [
        TargetReplica(sub_component_type=SubComponentType.PREFILL, desired_replicas=3)
    ]
    mock_kube_api.get_graph_deployment.side_effect = DynamoGraphDeploymentNotFoundError(
        "test-graph", "default"
    )

    # Act & Assert
    with pytest.raises(DynamoGraphDeploymentNotFoundError):
        await kubernetes_connector.set_component_replicas(target_replicas)


@pytest.mark.asyncio
async def test_set_component_replicas_empty_target_replicas(
    kubernetes_connector, mock_kube_api
):
    # Arrange
    target_replicas: list[TargetReplica] = []

    # Act & Assert
    with pytest.raises(EmptyTargetReplicasError):
        await kubernetes_connector.set_component_replicas(target_replicas)


async def test_set_component_replicas_deployment_not_ready(
    kubernetes_connector, mock_kube_api
):
    # Arrange
    target_replicas = [
        TargetReplica(sub_component_type=SubComponentType.PREFILL, desired_replicas=3),
        TargetReplica(sub_component_type=SubComponentType.DECODE, desired_replicas=2),
    ]
    mock_deployment = {
        "metadata": {"name": "test-graph"},
        "spec": {
            "services": {
                "component1": {"replicas": 1, "subComponentType": "prefill"},
                "component2": {"replicas": 2, "subComponentType": "decode"},
            }
        },
    }
    mock_kube_api.get_graph_deployment.return_value = mock_deployment
    mock_kube_api.is_deployment_ready.return_value = False

    # Act & Assert
    await kubernetes_connector.set_component_replicas(target_replicas)

    mock_kube_api.get_graph_deployment.assert_called_once()
    mock_kube_api.is_deployment_ready.assert_called_once_with(mock_deployment)
    mock_kube_api.update_graph_replicas.assert_not_called()
    mock_kube_api.wait_for_graph_deployment_ready.assert_not_called()


@pytest.mark.asyncio
async def test_validate_deployment_true(kubernetes_connector, mock_kube_api):
    # Arrange
    mock_deployment = {
        "metadata": {"name": "test-graph"},
        "spec": {
            "services": {
                "component1": {
                    "replicas": 1,
                    "subComponentType": "prefill",
                    "extraPodSpec": {
                        "mainContainer": {
                            "args": ["--served-model-name", "prefill-model"]
                        }
                    },
                },
                "component2": {"replicas": 2, "subComponentType": "decode"},
            }
        },
    }
    mock_kube_api.get_graph_deployment.return_value = mock_deployment

    # Act
    await kubernetes_connector.validate_deployment(decode_component_name="component2")


@pytest.mark.asyncio
async def test_validate_deployment_fail(kubernetes_connector, mock_kube_api):
    # Arrange
    mock_deployment = {
        "metadata": {"name": "test-graph"},
        "spec": {
            "services": {
                "component1": {"replicas": 1, "subComponentType": "prefill"},
                "component2": {"replicas": 2, "subComponentType": "prefill"},
            }
        },
    }
    mock_kube_api.get_graph_deployment.return_value = mock_deployment

    # Act
    with pytest.raises(DeploymentValidationError) as exc_info:
        await kubernetes_connector.validate_deployment()

    exception = exc_info.value
    assert set(exception.errors) == {
        str(DuplicateSubComponentError("prefill", ["component1", "component2"])),
        str(SubComponentNotFoundError("decode")),
    }


def test_get_model_name_both_none_raises_error(kubernetes_connector, mock_kube_api):
    # Arrange
    mock_deployment = {
        "metadata": {"name": "test-graph"},
        "spec": {
            "services": {
                "component1": {"replicas": 1, "subComponentType": "prefill"},
                "component2": {"replicas": 2, "subComponentType": "decode"},
            }
        },
    }

    with pytest.raises(ModelNameNotFoundError):
        kubernetes_connector.get_model_name(mock_deployment)


def test_get_model_name_prefill_none_decode_valid_returns_decode(kubernetes_connector):
    # Arrange
    mock_deployment = {
        "metadata": {"name": "test-graph"},
        "spec": {
            "services": {
                "component1": {"replicas": 1, "subComponentType": "prefill"},
                "component2": {
                    "replicas": 2,
                    "subComponentType": "decode",
                    "extraPodSpec": {
                        "mainContainer": {"args": ["--served-model-name", "test-model"]}
                    },
                },
            }
        },
    }
    # Act
    result = kubernetes_connector.get_model_name(mock_deployment)

    # Assert
    assert result == "test-model"


def test_get_model_name_mismatch_raises_error(kubernetes_connector, mock_kube_api):
    mock_deployment = {
        "metadata": {"name": "test-graph"},
        "spec": {
            "services": {
                "component1": {
                    "replicas": 1,
                    "subComponentType": "prefill",
                    "extraPodSpec": {
                        "mainContainer": {
                            "args": ["--served-model-name", "prefill-model"]
                        }
                    },
                },
                "component2": {
                    "replicas": 2,
                    "subComponentType": "decode",
                    "extraPodSpec": {
                        "mainContainer": {
                            "args": ["--served-model-name", "decode-model"]
                        }
                    },
                },
            }
        },
    }
    mock_kube_api.get_graph_deployment.return_value = mock_deployment

    # Act & Assert
    with pytest.raises(DeploymentModelNameMismatchError) as exc_info:
        kubernetes_connector.get_model_name(mock_deployment)

    exception = exc_info.value
    assert exception.prefill_model_name == "prefill-model"
    assert exception.decode_model_name == "decode-model"


def test_get_model_name_agree_returns_model_name(kubernetes_connector, mock_kube_api):
    # Arrange
    mock_deployment = {
        "metadata": {"name": "test-graph"},
        "spec": {
            "services": {
                "component1": {
                    "replicas": 1,
                    "subComponentType": "prefill",
                    "extraPodSpec": {
                        "mainContainer": {
                            "args": ["--served-model-name", "agreed-model"]
                        }
                    },
                },
                "component2": {
                    "replicas": 2,
                    "subComponentType": "decode",
                    "extraPodSpec": {
                        "mainContainer": {
                            "args": ["--served-model-name", "agreed-model"]
                        }
                    },
                },
            }
        },
    }
    mock_kube_api.get_graph_deployment.return_value = mock_deployment

    # Act
    result = kubernetes_connector.get_model_name(mock_deployment)

    # Assert
    assert result == "agreed-model"


# Tests for Service.get_gpu_count()
def test_service_get_gpu_count_valid():
    """Test that get_gpu_count returns correct GPU count from resources.limits.gpu"""
    service = Service(
        name="test-service",
        service={
            "replicas": 1,
            "resources": {"limits": {"gpu": "4"}},
        },
    )
    assert service.get_gpu_count() == 4


def test_service_get_gpu_count_from_requests_fallback():
    """Test that get_gpu_count falls back to requests.gpu when limits.gpu is missing"""
    service = Service(
        name="test-service",
        service={
            "replicas": 1,
            "resources": {"requests": {"gpu": "2"}},
        },
    )
    assert service.get_gpu_count() == 2


def test_service_get_gpu_count_limits_preferred_over_requests():
    """Test that limits.gpu is preferred over requests.gpu when both are present"""
    service = Service(
        name="test-service",
        service={
            "replicas": 1,
            "resources": {
                "limits": {"gpu": "4"},
                "requests": {"gpu": "2"},
            },
        },
    )
    assert service.get_gpu_count() == 4


def test_service_get_gpu_count_integer_value():
    """Test that get_gpu_count works with integer GPU values"""
    service = Service(
        name="test-service",
        service={
            "replicas": 1,
            "resources": {"limits": {"gpu": 2}},
        },
    )
    assert service.get_gpu_count() == 2


def test_service_get_gpu_count_missing_raises_error():
    """Test that get_gpu_count raises ValueError when GPU count is missing"""
    service = Service(
        name="test-service",
        service={"replicas": 1},
    )
    with pytest.raises(ValueError) as exc_info:
        service.get_gpu_count()
    assert "No GPU count specified" in str(exc_info.value)
    assert "test-service" in str(exc_info.value)


def test_service_get_gpu_count_invalid_raises_error():
    """Test that get_gpu_count raises ValueError for invalid GPU count"""
    service = Service(
        name="test-service",
        service={
            "replicas": 1,
            "resources": {"limits": {"gpu": "invalid"}},
        },
    )
    with pytest.raises(ValueError) as exc_info:
        service.get_gpu_count()
    assert "Invalid GPU count" in str(exc_info.value)


# Tests for KubernetesConnector.get_gpu_counts()
def test_get_gpu_counts_both_services(kubernetes_connector, mock_kube_api):
    """Test get_gpu_counts returns correct counts for both prefill and decode"""
    mock_deployment = {
        "metadata": {"name": "test-graph"},
        "spec": {
            "services": {
                "prefill-worker": {
                    "replicas": 1,
                    "subComponentType": "prefill",
                    "resources": {"limits": {"gpu": "2"}},
                },
                "decode-worker": {
                    "replicas": 1,
                    "subComponentType": "decode",
                    "resources": {"limits": {"gpu": "4"}},
                },
            }
        },
    }
    mock_kube_api.get_graph_deployment.return_value = mock_deployment

    prefill_gpu, decode_gpu = kubernetes_connector.get_gpu_counts()

    assert prefill_gpu == 2
    assert decode_gpu == 4


def test_get_gpu_counts_prefill_only(kubernetes_connector, mock_kube_api):
    """Test get_gpu_counts with require_decode=False"""
    mock_deployment = {
        "metadata": {"name": "test-graph"},
        "spec": {
            "services": {
                "prefill-worker": {
                    "replicas": 1,
                    "subComponentType": "prefill",
                    "resources": {"limits": {"gpu": "2"}},
                },
            }
        },
    }
    mock_kube_api.get_graph_deployment.return_value = mock_deployment

    prefill_gpu, decode_gpu = kubernetes_connector.get_gpu_counts(
        require_prefill=True, require_decode=False
    )

    assert prefill_gpu == 2
    assert decode_gpu == 0


def test_get_gpu_counts_decode_only(kubernetes_connector, mock_kube_api):
    """Test get_gpu_counts with require_prefill=False"""
    mock_deployment = {
        "metadata": {"name": "test-graph"},
        "spec": {
            "services": {
                "decode-worker": {
                    "replicas": 1,
                    "subComponentType": "decode",
                    "resources": {"limits": {"gpu": "4"}},
                },
            }
        },
    }
    mock_kube_api.get_graph_deployment.return_value = mock_deployment

    prefill_gpu, decode_gpu = kubernetes_connector.get_gpu_counts(
        require_prefill=False, require_decode=True
    )

    assert prefill_gpu == 0
    assert decode_gpu == 4


def test_get_gpu_counts_missing_gpu_raises_error(kubernetes_connector, mock_kube_api):
    """Test get_gpu_counts raises DeploymentValidationError when GPU count missing"""
    mock_deployment = {
        "metadata": {"name": "test-graph"},
        "spec": {
            "services": {
                "prefill-worker": {
                    "replicas": 1,
                    "subComponentType": "prefill",
                    # No resources.limits.gpu
                },
                "decode-worker": {
                    "replicas": 1,
                    "subComponentType": "decode",
                    "resources": {"limits": {"gpu": "4"}},
                },
            }
        },
    }
    mock_kube_api.get_graph_deployment.return_value = mock_deployment

    with pytest.raises(DeploymentValidationError) as exc_info:
        kubernetes_connector.get_gpu_counts()

    assert "prefill GPU count" in str(exc_info.value)


def test_get_gpu_counts_service_not_found_raises_error(
    kubernetes_connector, mock_kube_api
):
    """Test get_gpu_counts raises DeploymentValidationError when service not found"""
    mock_deployment = {
        "metadata": {"name": "test-graph"},
        "spec": {
            "services": {
                "prefill-worker": {
                    "replicas": 1,
                    "subComponentType": "prefill",
                    "resources": {"limits": {"gpu": "2"}},
                },
                # No decode service
            }
        },
    }
    mock_kube_api.get_graph_deployment.return_value = mock_deployment

    with pytest.raises(DeploymentValidationError) as exc_info:
        kubernetes_connector.get_gpu_counts()

    assert "decode GPU count" in str(exc_info.value)


# Tests for get_actual_worker_counts


def test_get_actual_worker_counts_stable(kubernetes_connector, mock_kube_api):
    """Test get_actual_worker_counts when both services are stable"""
    mock_deployment = {
        "metadata": {"name": "test-graph"},
        "spec": {
            "services": {
                "prefill-component": {},
                "decode-component": {},
            }
        },
    }
    mock_kube_api.get_graph_deployment.return_value = mock_deployment
    mock_kube_api.get_service_replica_status.side_effect = [(2, True), (4, True)]

    (
        prefill_count,
        decode_count,
        is_stable,
    ) = kubernetes_connector.get_actual_worker_counts(
        prefill_component_name="prefill-component",
        decode_component_name="decode-component",
    )

    assert prefill_count == 2
    assert decode_count == 4
    assert is_stable is True


def test_get_actual_worker_counts_prefill_rollout_in_progress(
    kubernetes_connector, mock_kube_api
):
    """Test get_actual_worker_counts when prefill has rollout in progress"""
    mock_deployment = {
        "metadata": {"name": "test-graph"},
        "spec": {
            "services": {
                "prefill-component": {},
                "decode-component": {},
            }
        },
    }
    mock_kube_api.get_graph_deployment.return_value = mock_deployment
    mock_kube_api.get_service_replica_status.side_effect = [(2, False), (4, True)]

    (
        prefill_count,
        decode_count,
        is_stable,
    ) = kubernetes_connector.get_actual_worker_counts(
        prefill_component_name="prefill-component",
        decode_component_name="decode-component",
    )

    assert prefill_count == 2
    assert decode_count == 4
    assert is_stable is False


def test_get_actual_worker_counts_prefill_only(kubernetes_connector, mock_kube_api):
    """Test get_actual_worker_counts with only prefill component"""
    mock_deployment = {
        "metadata": {"name": "test-graph"},
        "spec": {
            "services": {
                "prefill-component": {
                    "replicas": 2,
                    "subComponentType": "prefill",
                },
            }
        },
    }
    mock_kube_api.get_graph_deployment.return_value = mock_deployment
    mock_kube_api.get_service_replica_status.return_value = (2, True)

    (
        prefill_count,
        decode_count,
        is_stable,
    ) = kubernetes_connector.get_actual_worker_counts(
        prefill_component_name="prefill-component",
        decode_component_name=None,
    )

    assert prefill_count == 2
    assert decode_count == 0
    assert is_stable is True


def test_get_actual_worker_counts_decode_only(kubernetes_connector, mock_kube_api):
    """Test get_actual_worker_counts with only decode component"""
    mock_deployment = {
        "metadata": {"name": "test-graph"},
        "spec": {
            "services": {
                "decode-component": {
                    "replicas": 4,
                    "subComponentType": "decode",
                },
            }
        },
    }
    mock_kube_api.get_graph_deployment.return_value = mock_deployment
    mock_kube_api.get_service_replica_status.return_value = (4, True)

    (
        prefill_count,
        decode_count,
        is_stable,
    ) = kubernetes_connector.get_actual_worker_counts(
        prefill_component_name=None,
        decode_component_name="decode-component",
    )

    assert prefill_count == 0
    assert decode_count == 4
    assert is_stable is True


def test_get_actual_worker_counts_no_components(kubernetes_connector, mock_kube_api):
    """Test get_actual_worker_counts with no components specified"""
    mock_deployment = {
        "metadata": {"name": "test-graph"},
        "spec": {"services": {}},
        "status": {"services": {}},
    }
    mock_kube_api.get_graph_deployment.return_value = mock_deployment

    (
        prefill_count,
        decode_count,
        is_stable,
    ) = kubernetes_connector.get_actual_worker_counts(
        prefill_component_name=None,
        decode_component_name=None,
    )

    assert prefill_count == 0
    assert decode_count == 0
    assert is_stable is True


# Tests for _resolve_dgd_service / get_worker_info component-filter.
#
# Regression: the filter that compares an MDC entry's ``component`` field
# against ``expected_component`` must use the lowercase backend-default
# name (what the Rust runtime writes to MDC), NOT the DGD ``spec.services``
# dict key. The DGD key is typically PascalCase (``VllmPrefillWorker``)
# while MDC carries the Endpoint name (``prefill`` / ``backend`` /
# ``tensorrt_llm``); returning the DGD key for the filter would cause
# every real-world MDC entry to be skipped, leaving WorkerInfo without
# ``context_length`` and silently breaking easy-mode load scaling.


def test_resolve_dgd_service_prefill_uses_backend_default_for_filter(
    kubernetes_connector, mock_kube_api
):
    """vLLM prefill: filter name = "prefill" (MDC side), not DGD services key."""
    mock_deployment = {
        "metadata": {"name": "test-graph"},
        "spec": {
            "services": {
                "VllmPrefillWorker": {
                    "replicas": 1,
                    "subComponentType": "prefill",
                },
            }
        },
    }
    mock_kube_api.get_graph_deployment.return_value = mock_deployment

    dgd_service_name, expected_component = kubernetes_connector._resolve_dgd_service(
        SubComponentType.PREFILL, backend="vllm"
    )

    # k8s operations (e.g. replica patch) still target the PascalCase DGD key.
    assert dgd_service_name == "VllmPrefillWorker"
    # The filter side must match what the Rust runtime writes to MDC.
    assert expected_component == "prefill"


def test_resolve_dgd_service_decode_uses_backend_default_for_filter(
    kubernetes_connector, mock_kube_api
):
    """vLLM decode: MDC carries "backend", NOT "decode"; filter must match that."""
    mock_deployment = {
        "metadata": {"name": "test-graph"},
        "spec": {
            "services": {
                "VllmDecodeWorker": {
                    "replicas": 1,
                    "subComponentType": "decode",
                },
            }
        },
    }
    mock_kube_api.get_graph_deployment.return_value = mock_deployment

    dgd_service_name, expected_component = kubernetes_connector._resolve_dgd_service(
        SubComponentType.DECODE, backend="vllm"
    )

    assert dgd_service_name == "VllmDecodeWorker"
    # Critically, vLLM's decode-worker component name is "backend" (from
    # VllmComponentName.decode_worker_component_name). Using
    # SubComponentType.DECODE.value ("decode") here would break decode
    # filtering on every backend.
    assert expected_component == "backend"


def test_resolve_dgd_service_trtllm_decode_uses_tensorrt_llm_name(
    kubernetes_connector, mock_kube_api
):
    """TRT-LLM decode: MDC carries "tensorrt_llm"; filter must match."""
    mock_deployment = {
        "metadata": {"name": "test-graph"},
        "spec": {
            "services": {
                "TRTLLMDecodeWorker": {
                    "replicas": 1,
                    "subComponentType": "decode",
                },
            }
        },
    }
    mock_kube_api.get_graph_deployment.return_value = mock_deployment

    _, expected_component = kubernetes_connector._resolve_dgd_service(
        SubComponentType.DECODE, backend="trtllm"
    )

    assert expected_component == "tensorrt_llm"


def test_resolve_dgd_service_missing_dgd_still_returns_backend_default(
    kubernetes_connector, mock_kube_api
):
    """When DGD lookup fails, still return the backend default for filtering."""
    mock_kube_api.get_graph_deployment.side_effect = DynamoGraphDeploymentNotFoundError(
        "test-graph", "default"
    )

    dgd_service_name, expected_component = kubernetes_connector._resolve_dgd_service(
        SubComponentType.PREFILL, backend="vllm"
    )

    assert dgd_service_name is None
    assert expected_component == "prefill"


def test_resolve_dgd_service_respects_user_endpoint_override(
    kubernetes_connector, mock_kube_api
):
    """If the DGD passes --endpoint ns.comp.ep, the MDC filter must use 'comp'."""
    mock_deployment = {
        "metadata": {"name": "test-graph"},
        "spec": {
            "services": {
                "VllmPrefillWorker": {
                    "replicas": 1,
                    "subComponentType": "prefill",
                    "extraPodSpec": {
                        "mainContainer": {
                            "args": [
                                "--endpoint",
                                "my-ns.my-custom-prefill.generate",
                                "--model",
                                "Qwen/Qwen3-8B",
                            ]
                        }
                    },
                },
            }
        },
    }
    mock_kube_api.get_graph_deployment.return_value = mock_deployment

    dgd_service_name, expected_component = kubernetes_connector._resolve_dgd_service(
        SubComponentType.PREFILL, backend="vllm"
    )

    # k8s operations still target the DGD services key.
    assert dgd_service_name == "VllmPrefillWorker"
    # Filter must match what the worker will actually write to MDC, which
    # comes from the user's --endpoint override, not the backend default.
    assert expected_component == "my-custom-prefill"


def test_resolve_dgd_service_endpoint_override_with_dyn_prefix(
    kubernetes_connector, mock_kube_api
):
    """parse_endpoint accepts 'dyn://' prefix; the extracted component must strip it."""
    mock_deployment = {
        "metadata": {"name": "test-graph"},
        "spec": {
            "services": {
                "VllmDecodeWorker": {
                    "replicas": 1,
                    "subComponentType": "decode",
                    "extraPodSpec": {
                        "mainContainer": {
                            "args": [
                                "--endpoint",
                                "dyn://ns.user-decode.generate",
                            ]
                        }
                    },
                },
            }
        },
    }
    mock_kube_api.get_graph_deployment.return_value = mock_deployment

    _, expected_component = kubernetes_connector._resolve_dgd_service(
        SubComponentType.DECODE, backend="vllm"
    )

    assert expected_component == "user-decode"


def test_resolve_dgd_service_malformed_endpoint_falls_back_to_default(
    kubernetes_connector, mock_kube_api
):
    """Malformed --endpoint (wrong number of parts) falls back to backend default."""
    mock_deployment = {
        "metadata": {"name": "test-graph"},
        "spec": {
            "services": {
                "VllmPrefillWorker": {
                    "replicas": 1,
                    "subComponentType": "prefill",
                    "extraPodSpec": {
                        "mainContainer": {
                            "args": ["--endpoint", "only-two.parts"],
                        }
                    },
                },
            }
        },
    }
    mock_kube_api.get_graph_deployment.return_value = mock_deployment

    _, expected_component = kubernetes_connector._resolve_dgd_service(
        SubComponentType.PREFILL, backend="vllm"
    )

    assert expected_component == "prefill"


def test_service_get_component_name_from_endpoint_arg_present():
    service = Service(
        name="VllmPrefillWorker",
        service={
            "extraPodSpec": {
                "mainContainer": {
                    "args": [
                        "--endpoint",
                        "ns.custom-comp.generate",
                        "--other",
                        "flag",
                    ]
                }
            }
        },
    )
    assert service.get_component_name_from_endpoint_arg() == "custom-comp"


def test_service_get_component_name_from_endpoint_arg_absent():
    service = Service(
        name="VllmPrefillWorker",
        service={
            "extraPodSpec": {
                "mainContainer": {
                    "args": ["--model", "Qwen/Qwen3-8B"],
                }
            }
        },
    )
    assert service.get_component_name_from_endpoint_arg() is None


def test_service_get_component_name_from_endpoint_arg_missing_value():
    """--endpoint with no following arg should return None, not raise IndexError."""
    service = Service(
        name="VllmPrefillWorker",
        service={"extraPodSpec": {"mainContainer": {"args": ["--endpoint"]}}},
    )
    assert service.get_component_name_from_endpoint_arg() is None
