#!/usr/bin/env python3
# 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.
"""
Test script for v1beta1 Pydantic models.

Validates that the generated Pydantic models can be imported and used correctly.
"""

import subprocess
import sys
import types
from pathlib import Path


def _repo_root() -> Path:
    start = Path(__file__).parent
    try:
        result = subprocess.run(
            ["git", "rev-parse", "--show-toplevel"],
            capture_output=True,
            text=True,
            check=True,
            cwd=start,
        )
        return Path(result.stdout.strip())
    except (subprocess.CalledProcessError, FileNotFoundError):
        pass
    # Fallback: walk up until we find go.mod (same logic as generate_pydantic_from_go.py)
    p = start
    while p != p.parent:
        if (p / "go.mod").exists():
            return p
        p = p.parent
    return start


_components_src = _repo_root() / "components" / "src"

# In the operator Docker build the context is deploy/operator/ only — components/src
# is not copied in. The generated files are already committed, so skip validation.
if not _components_src.exists():
    print(
        f"Note: {_components_src} not found (operator-only build context). "
        "Skipping Pydantic validation tests."
    )
    sys.exit(0)

# Add the components src to path so we can import the generated models
sys.path.insert(0, str(_components_src))

# ---------------------------------------------------------------------------
# Stub dynamo.runtime.logging and bypass the heavy dynamo.planner.__init__
# before importing any dynamo module.
#
# dynamo itself must be a namespace-like package (has __path__) so that
# Python's import machinery can traverse down to dynamo.profiler from the
# filesystem.  dynamo.planner is pre-registered as a stub to skip its heavy
# __init__.py, while still allowing dynamo.planner.config.* to load normally.
# ---------------------------------------------------------------------------
_dynamo_path = str(_components_src / "dynamo")
_planner_path = str(_components_src / "dynamo" / "planner")

if "dynamo" not in sys.modules:
    _dynamo_mod = types.ModuleType("dynamo")
    _dynamo_mod.__path__ = [_dynamo_path]  # type: ignore[attr-defined]
    _dynamo_mod.__package__ = "dynamo"
    sys.modules["dynamo"] = _dynamo_mod

if "dynamo.runtime" not in sys.modules:
    _runtime_mod = types.ModuleType("dynamo.runtime")
    sys.modules["dynamo.runtime"] = _runtime_mod

_logging_mod = types.ModuleType("dynamo.runtime.logging")
_logging_mod.configure_dynamo_logging = lambda *args, **kwargs: None  # type: ignore[attr-defined]
sys.modules["dynamo.runtime.logging"] = _logging_mod

_planner_mod = types.ModuleType("dynamo.planner")
_planner_mod.__path__ = [_planner_path]  # type: ignore[attr-defined]
_planner_mod.__package__ = "dynamo.planner"
sys.modules["dynamo.planner"] = _planner_mod

import pydantic  # noqa: E402

from dynamo.profiler.utils.dgdr_v1beta1_types import (  # noqa: E402
    BackendType,
    DeploymentInfoStatus,
    DGDRPhase,
    DynamoGraphDeploymentRequestSpec,
    DynamoGraphDeploymentRequestStatus,
    FeaturesSpec,
    MockerSpec,
    ModelCacheSpec,
    PlannerConfig,
    PlannerPreDeploymentSweepMode,
    ProfilingPhase,
    SearchStrategy,
    SLASpec,
    WorkloadSpec,
)

print("✓ Successfully imported all Pydantic models")


def test_simple_dgdr():
    """Test creating a simple DGDR (minimal spec)"""
    spec = DynamoGraphDeploymentRequestSpec(
        model="Qwen/Qwen3-32B",
    )
    print("✓ Created simple DGDR spec")

    assert spec.model == "Qwen/Qwen3-32B"
    assert spec.backend == BackendType.Auto  # kubebuilder:default=auto
    assert spec.autoApply is True  # kubebuilder:default=true
    print("✓ Simple DGDR spec validation passed")


def test_full_dgdr():
    """Test creating a full DGDR with all fields"""
    spec = DynamoGraphDeploymentRequestSpec(
        model="meta-llama/Llama-3.1-405B",
        backend=BackendType.Vllm,
        image="nvcr.io/nvidia/dynamo-runtime:latest",
        workload=WorkloadSpec(
            isl=1024,
            osl=512,
            concurrency=10.0,
        ),
        sla=SLASpec(
            ttft=100.0,
            itl=10.0,
        ),
        modelCache=ModelCacheSpec(
            pvcName="model-cache",
            pvcModelPath="llama-3.1-405b",
        ),
        features=FeaturesSpec(
            planner=PlannerConfig(enable_load_scaling=False),
            mocker=MockerSpec(enabled=False),
        ),
        searchStrategy=SearchStrategy.Rapid,
        autoApply=True,
    )
    print("✓ Created full DGDR spec")

    assert spec.model == "meta-llama/Llama-3.1-405B"
    assert spec.backend == BackendType.Vllm
    assert spec.workload.isl == 1024
    assert spec.sla.ttft == 100.0
    assert spec.sla.itl == 10.0
    assert spec.modelCache.pvcName == "model-cache"
    assert spec.modelCache.pvcModelPath == "llama-3.1-405b"
    assert isinstance(spec.features.planner, PlannerConfig)
    assert spec.features.mocker.enabled is False
    print("✓ Full DGDR spec validation passed")


def test_sla_defaults_and_validation():
    """Test SLASpec defaults and mutual-exclusivity validator"""
    # Default mode: ttft + itl with python-defaults
    sla = SLASpec()
    assert sla.ttft == 2000.0
    assert sla.itl == 30.0
    assert sla.e2eLatency is None
    print("✓ SLASpec defaults correct")

    # explicit ttft+itl mode: OK
    SLASpec(ttft=100.0, itl=10.0)

    # e2eLatency mode: OK (null out ttft/itl)
    SLASpec(ttft=None, itl=None, e2eLatency=500.0)

    # e2eLatency mode: OK without explicitly nulling defaults
    SLASpec(e2eLatency=500.0)

    # mixing modes should raise
    try:
        SLASpec(ttft=100.0, itl=10.0, e2eLatency=500.0)
        raise AssertionError("expected ValidationError for mixed SLA modes")
    except pydantic.ValidationError:
        pass

    # ttft without itl should raise
    try:
        SLASpec(itl=None, ttft=100.0)
        raise AssertionError("expected ValidationError for ttft without itl")
    except pydantic.ValidationError:
        pass

    print("✓ SLASpec validation correct")


def test_workload_defaults():
    """Test WorkloadSpec kubebuilder defaults"""
    w = WorkloadSpec()
    assert w.isl == 4000
    assert w.osl == 1000
    print("✓ WorkloadSpec defaults correct")


def test_enums():
    """Test enum values"""
    # DGDRPhase — TitleCase suffix from Go const names
    assert DGDRPhase.Pending == "Pending"
    assert DGDRPhase.Profiling == "Profiling"
    assert DGDRPhase.Deployed == "Deployed"

    # ProfilingPhase — TitleCase suffix from Go const names
    assert ProfilingPhase.Initializing == "Initializing"
    assert ProfilingPhase.SweepingPrefill == "SweepingPrefill"

    # SearchStrategy — TitleCase from Go const names
    assert SearchStrategy.Rapid == "rapid"
    assert SearchStrategy.Thorough == "thorough"

    # BackendType — mixed case from Go const names
    assert BackendType.Auto == "auto"
    assert BackendType.Vllm == "vllm"

    # PlannerPreDeploymentSweepMode (None → None_ to avoid Python keyword clash)
    assert PlannerPreDeploymentSweepMode.None_ == "none"
    assert PlannerPreDeploymentSweepMode.Rapid == "rapid"

    print("✓ All enum values validated")


def test_status_models():
    """Test status model creation"""
    status = DynamoGraphDeploymentRequestStatus(
        phase=DGDRPhase.Profiling,
        profilingPhase=ProfilingPhase.SweepingPrefill,
        dgdName="test-dgd",
        profilingJobName="test-profiling-job",
        deploymentInfo=DeploymentInfoStatus(
            replicas=3,
            availableReplicas=2,
        ),
    )
    print("✓ Created DGDR status")

    assert status.phase == DGDRPhase.Profiling
    assert status.profilingPhase == ProfilingPhase.SweepingPrefill
    assert status.deploymentInfo.replicas == 3
    print("✓ DGDR status validation passed")


def main():
    """Run all tests"""
    print("\n" + "=" * 60)
    print("Testing v1beta1 Pydantic Models")
    print("=" * 60 + "\n")

    test_simple_dgdr()
    test_full_dgdr()
    test_sla_defaults_and_validation()
    test_workload_defaults()
    test_enums()
    test_status_models()

    print("\n" + "=" * 60)
    print("All tests passed! ✓")
    print("=" * 60 + "\n")

    return 0


if __name__ == "__main__":
    sys.exit(main())