"""
Functional test based on Vajra's configuration hierarchy.

This test validates that vidhi can handle a real-world complex config structure
with multiple polymorphic hierarchies, deep nesting, and CLI argument generation.

Patterns covered from vajra:
- Polymorphic config hierarchies with enum type discriminators
- Native handle pattern (simulated without C++ bindings)
- Computed fields set in __post_init__
- Cross-field validation (min < max, divisibility checks)
- Properties on frozen dataclasses
- Optional fields with lambda factories
- File path templates with placeholder substitution
- Methods that compute values from nested configs
- Conditional field modification in __post_init__
- Range and positive value validation
"""

import os
import sys
import tempfile
from dataclasses import field as dataclass_field
from enum import Enum
from io import StringIO
from pathlib import Path
from typing import List, Optional
from unittest.mock import patch

import pytest
import yaml

from vidhi import (
    BasePolyConfig,
    create_class_from_dict,
    dataclass_to_dict,
    field,
    frozen_dataclass,
    load_yaml_config,
    parse_cli_args,
)
from vidhi.nested_config import ConfigWalker


# =============================================================================
# Enums (mirroring vajra's C++ enums)
# =============================================================================
class AllocatorType(Enum):
    LRU = "lru"
    LFU = "lfu"
    COST_AWARE_EXP_DECAY = "cost_aware_exp_decay"
    COST_AWARE_GDSF = "cost_aware_gdsf"


class ComputeManagerType(Enum):
    FIXED_CHUNK = "fixed_chunk"
    DYNAMIC_CHUNK = "dynamic_chunk"
    SPACE_SHARING = "space_sharing"


class SessionRouterType(Enum):
    PULL = "pull"
    ROUND_ROBIN = "round_robin"


class SessionPrioritizerType(Enum):
    FCFS = "fcfs"
    EDF = "edf"
    LRS = "lrs"


class ResourceAllocatorType(Enum):
    LOCAL = "local"
    RAY = "ray"
    MOCK = "mock"


class ModelConfigType(Enum):
    LLM = "llm"


class ReplicaControllerType(Enum):
    LLM_BASE = "llm_base"


class ReplicasetControllerType(Enum):
    LLM = "llm"


class SchedulerType(Enum):
    GREEDY = "greedy"
    BALANCED = "balanced"
    PRIORITY = "priority"


class StorageBackendType(Enum):
    LOCAL = "local"
    S3 = "s3"
    GCS = "gcs"


# Non-polymorphic enum (mirroring vajra's pybind11 AttentionBackendType)
class AttentionBackendType(Enum):
    FLASH_ATTENTION = "flash_attention"
    FLASHINFER = "flashinfer"


class ProfilingScheduleType(Enum):
    BATCHWISE = "batchwise"
    CYCLICAL = "cyclical"


# =============================================================================
# Polymorphic Hierarchy 1: Allocator (4 variants)
# =============================================================================
@frozen_dataclass
class AbstractAllocatorConfig(BasePolyConfig):
    """Abstract base for cache allocator configurations."""

    @classmethod
    def get_type(cls) -> AllocatorType:
        raise NotImplementedError()


@frozen_dataclass
class LruAllocatorConfig(AbstractAllocatorConfig):
    """LRU (Least Recently Used) allocator."""

    @staticmethod
    def get_type() -> AllocatorType:
        return AllocatorType.LRU


@frozen_dataclass
class LfuAllocatorConfig(AbstractAllocatorConfig):
    """LFU (Least Frequently Used) allocator."""

    @staticmethod
    def get_type() -> AllocatorType:
        return AllocatorType.LFU


@frozen_dataclass
class CostAwareExpDecayAllocatorConfig(AbstractAllocatorConfig):
    """Cost-aware allocator with exponential decay."""

    prefill_profiling_path: str = field(
        "./profiling.json", help="Path to prefill profiling JSON"
    )
    decay_rate: float = field(0.5, help="Forward decay rate")
    min_cost: float = field(1e-10, help="Minimum cost value")

    @staticmethod
    def get_type() -> AllocatorType:
        return AllocatorType.COST_AWARE_EXP_DECAY

    def __post_init__(self):
        if self.decay_rate <= 0:
            raise ValueError("decay_rate must be positive")
        if self.min_cost <= 0:
            raise ValueError("min_cost must be positive")


@frozen_dataclass
class CostAwareGdsfAllocatorConfig(AbstractAllocatorConfig):
    """Cost-aware allocator with GDSF aging."""

    prefill_profiling_path: str = field(
        "./profiling.json", help="Path to prefill profiling JSON"
    )
    use_frequency: bool = field(True, help="Use GDSF (with frequency) vs GDS")

    @staticmethod
    def get_type() -> AllocatorType:
        return AllocatorType.COST_AWARE_GDSF


# =============================================================================
# Polymorphic Hierarchy 2: Compute Manager (3 variants)
# =============================================================================
@frozen_dataclass
class AbstractComputeManagerConfig(BasePolyConfig):
    """Abstract base for compute manager configurations."""

    @classmethod
    def get_type(cls) -> ComputeManagerType:
        raise NotImplementedError()


@frozen_dataclass
class FixedChunkComputeManagerConfig(AbstractComputeManagerConfig):
    """Fixed chunk compute manager."""

    chunk_size: int = field(2048, help="Fixed chunk size in tokens")

    @staticmethod
    def get_type() -> ComputeManagerType:
        return ComputeManagerType.FIXED_CHUNK


@frozen_dataclass
class DynamicChunkComputeManagerConfig(AbstractComputeManagerConfig):
    """Dynamic chunk compute manager with cross-field validation."""

    min_chunk_size: int = field(512, help="Minimum chunk size")
    max_chunk_size: int = field(4096, help="Maximum chunk size")

    @staticmethod
    def get_type() -> ComputeManagerType:
        return ComputeManagerType.DYNAMIC_CHUNK

    def __post_init__(self):
        # Cross-field validation (vajra pattern: min < max)
        if self.min_chunk_size <= 0:
            raise ValueError("min_chunk_size must be positive")
        if self.max_chunk_size <= 0:
            raise ValueError("max_chunk_size must be positive")
        if self.min_chunk_size >= self.max_chunk_size:
            raise ValueError("min_chunk_size must be less than max_chunk_size")


@frozen_dataclass
class SpaceSharingComputeManagerConfig(AbstractComputeManagerConfig):
    """Space sharing compute manager."""

    max_concurrent_sequences: int = field(8, help="Max concurrent sequences")

    @staticmethod
    def get_type() -> ComputeManagerType:
        return ComputeManagerType.SPACE_SHARING


# =============================================================================
# Polymorphic Hierarchy 3: Session Router (2 variants)
# =============================================================================
@frozen_dataclass
class AbstractSessionRouterConfig(BasePolyConfig):
    """Abstract base for session router configurations."""

    @classmethod
    def get_type(cls) -> SessionRouterType:
        raise NotImplementedError()


@frozen_dataclass
class PullSessionRouterConfig(AbstractSessionRouterConfig):
    """Pull-based session router."""

    @staticmethod
    def get_type() -> SessionRouterType:
        return SessionRouterType.PULL


@frozen_dataclass
class RoundRobinSessionRouterConfig(AbstractSessionRouterConfig):
    """Round-robin session router."""

    @staticmethod
    def get_type() -> SessionRouterType:
        return SessionRouterType.ROUND_ROBIN


# =============================================================================
# Polymorphic Hierarchy 4: Session Prioritizer (3 variants)
# =============================================================================
@frozen_dataclass
class AbstractSessionPrioritizerConfig(BasePolyConfig):
    """Abstract base for session prioritizer configurations."""

    @classmethod
    def get_type(cls) -> SessionPrioritizerType:
        raise NotImplementedError()


@frozen_dataclass
class FcfsSessionPrioritizerConfig(AbstractSessionPrioritizerConfig):
    """First-come-first-served prioritizer."""

    @staticmethod
    def get_type() -> SessionPrioritizerType:
        return SessionPrioritizerType.FCFS


@frozen_dataclass
class EdfSessionPrioritizerConfig(AbstractSessionPrioritizerConfig):
    """Earliest deadline first prioritizer."""

    default_deadline_ms: float = field(1000.0, help="Default deadline in milliseconds")

    @staticmethod
    def get_type() -> SessionPrioritizerType:
        return SessionPrioritizerType.EDF


@frozen_dataclass
class LrsSessionPrioritizerConfig(AbstractSessionPrioritizerConfig):
    """Longest remaining service prioritizer."""

    preemption_threshold: float = field(0.5, help="Preemption threshold")

    @staticmethod
    def get_type() -> SessionPrioritizerType:
        return SessionPrioritizerType.LRS


# =============================================================================
# Polymorphic Hierarchy 5: Resource Allocator (3 variants)
# =============================================================================
@frozen_dataclass
class AbstractResourceAllocatorConfig(BasePolyConfig):
    """Abstract base for resource allocator configurations."""

    @classmethod
    def get_type(cls) -> ResourceAllocatorType:
        raise NotImplementedError()


@frozen_dataclass
class LocalGpuResourceAllocatorConfig(AbstractResourceAllocatorConfig):
    """Local GPU resource allocator."""

    @staticmethod
    def get_type() -> ResourceAllocatorType:
        return ResourceAllocatorType.LOCAL


@frozen_dataclass
class RayGpuResourceAllocatorConfig(AbstractResourceAllocatorConfig):
    """Ray-based GPU resource allocator."""

    ray_address: str = field("auto", help="Ray cluster address")

    @staticmethod
    def get_type() -> ResourceAllocatorType:
        return ResourceAllocatorType.RAY


@frozen_dataclass
class MockResourceAllocatorConfig(AbstractResourceAllocatorConfig):
    """Mock resource allocator for testing."""

    num_mock_gpus: int = field(4, help="Number of mock GPUs")

    @staticmethod
    def get_type() -> ResourceAllocatorType:
        return ResourceAllocatorType.MOCK


# =============================================================================
# Regular (non-polymorphic) nested configs
# =============================================================================
@frozen_dataclass
class CacheTierConfig:
    """Configuration for a single cache tier."""

    admission_watermark: float = field(0.05, help="Admission watermark threshold")
    demotion_watermark: float = field(0.1, help="Demotion watermark threshold")

    def __post_init__(self):
        if not 0 <= self.admission_watermark <= 1:
            raise ValueError("admission_watermark must be between 0 and 1")
        if not 0 <= self.demotion_watermark <= 1:
            raise ValueError("demotion_watermark must be between 0 and 1")
        # Cross-field validation (vajra pattern)
        if self.admission_watermark >= self.demotion_watermark:
            raise ValueError("admission_watermark must be less than demotion_watermark")


@frozen_dataclass
class CacheConfig:
    """Multi-tiered cache configuration."""

    page_size: int = field(16, help="Cache page size in tokens")
    max_promotion_batch_size: int = field(1000, help="Max tokens to promote per batch")
    max_demotion_batch_size: int = field(1000, help="Max tokens to demote per batch")
    allocator: AbstractAllocatorConfig = dataclass_field(
        default_factory=LruAllocatorConfig
    )
    gpu_tier: CacheTierConfig = dataclass_field(default_factory=CacheTierConfig)
    cpu_tier: CacheTierConfig = dataclass_field(default_factory=CacheTierConfig)
    nvme_tier: CacheTierConfig = dataclass_field(default_factory=CacheTierConfig)
    enable_prefix_caching: bool = field(True, help="Enable prefix caching")

    def __post_init__(self):
        if self.page_size <= 0:
            raise ValueError("page_size must be positive")


@frozen_dataclass
class WorkerConfig:
    """Worker process configuration."""

    num_workers: int = field(1, help="Number of worker processes")
    worker_use_ray: bool = field(False, help="Use Ray for workers")
    attention_backend_type: AttentionBackendType = field(
        AttentionBackendType.FLASH_ATTENTION,
        help="Attention backend: FLASH_ATTENTION or FLASHINFER.",
    )


# =============================================================================
# Polymorphic Hierarchy 6: Model Deployment
# =============================================================================
@frozen_dataclass
class AbstractModelDeploymentConfig(BasePolyConfig):
    """Abstract base for model deployment configurations."""

    model: str = field("meta-llama/Llama-3-8B", help="Model name or path")
    tensor_parallel_size: int = field(1, help="Tensor parallel size")
    pipeline_parallel_size: int = field(1, help="Pipeline parallel size")
    max_model_len: int = field(4096, help="Maximum model length")
    seed: int = field(0, help="Random seed")

    @property
    def world_size(self) -> int:
        return self.tensor_parallel_size * self.pipeline_parallel_size

    @classmethod
    def get_type(cls) -> ModelConfigType:
        raise NotImplementedError()


@frozen_dataclass
class ModelDeploymentConfig(AbstractModelDeploymentConfig):
    """LLM model deployment configuration."""

    trust_remote_code: bool = field(True, help="Trust remote code")
    load_format: str = field("safetensors", help="Model weights format")

    @staticmethod
    def get_type() -> ModelConfigType:
        return ModelConfigType.LLM

    def __post_init__(self):
        if self.load_format not in ["safetensors", "dummy"]:
            raise ValueError(f"Invalid load_format: {self.load_format}")


# =============================================================================
# Polymorphic Hierarchy 7: Replica Controller (2 levels deep)
# =============================================================================
@frozen_dataclass
class AbstractReplicaControllerConfig(BasePolyConfig):
    """Abstract base for replica controller configurations."""

    @classmethod
    def get_type(cls) -> ReplicaControllerType:
        raise NotImplementedError()


def _get_default_model_deployment():
    """Lazy import to avoid circular dependencies."""
    return ModelDeploymentConfig()


@frozen_dataclass
class ReplicaControllerConfig(AbstractReplicaControllerConfig):
    """LLM replica controller configuration."""

    model_deployment_config: AbstractModelDeploymentConfig = dataclass_field(
        default_factory=_get_default_model_deployment
    )
    worker_config: WorkerConfig = dataclass_field(default_factory=WorkerConfig)
    cache_config: CacheConfig = dataclass_field(default_factory=CacheConfig)
    compute_manager_config: AbstractComputeManagerConfig = dataclass_field(
        default_factory=FixedChunkComputeManagerConfig
    )

    @staticmethod
    def get_type() -> ReplicaControllerType:
        return ReplicaControllerType.LLM_BASE


# =============================================================================
# Polymorphic Hierarchy 8: Replicaset Controller (3 levels deep)
# =============================================================================
@frozen_dataclass
class AbstractReplicasetControllerConfig(BasePolyConfig):
    """Abstract base for replicaset controller configurations."""

    @classmethod
    def get_type(cls) -> ReplicasetControllerType:
        raise NotImplementedError()


def _get_default_replica_controller():
    return ReplicaControllerConfig()


@frozen_dataclass
class ReplicasetControllerConfig(AbstractReplicasetControllerConfig):
    """LLM replicaset controller configuration."""

    replica_controller_config: AbstractReplicaControllerConfig = dataclass_field(
        default_factory=_get_default_replica_controller
    )
    session_prioritizer_config: AbstractSessionPrioritizerConfig = dataclass_field(
        default_factory=FcfsSessionPrioritizerConfig
    )
    session_router_config: AbstractSessionRouterConfig = dataclass_field(
        default_factory=PullSessionRouterConfig
    )

    @staticmethod
    def get_type() -> ReplicasetControllerType:
        return ReplicasetControllerType.LLM


# =============================================================================
# Top-level config: InferenceEngineConfig
# =============================================================================
def _get_default_replicaset_controller():
    return ReplicasetControllerConfig()


# =============================================================================
# Polymorphic Hierarchy 9: Profiling Schedule (deeply nested in MetricsConfig)
# This tests the pattern where a non-polymorphic config contains a deeply
# nested polymorphic field — the bug that caused MetricsConfig fields to
# disappear from CLI help.
# =============================================================================
@frozen_dataclass
class AbstractProfilingSchedule(BasePolyConfig):
    """Abstract base for profiling schedule configurations."""

    @classmethod
    def get_type(cls) -> ProfilingScheduleType:
        raise NotImplementedError()


@frozen_dataclass
class BatchwiseProfilingSchedule(AbstractProfilingSchedule):
    """Batchwise profiling schedule."""

    traces_per_batch_shape: int = field(3, help="Max traces per batch hash")

    @staticmethod
    def get_type() -> ProfilingScheduleType:
        return ProfilingScheduleType.BATCHWISE


@frozen_dataclass
class CyclicalProfilingSchedule(AbstractProfilingSchedule):
    """Cyclical profiling schedule."""

    warmup_batches: int = field(0, help="Batches to skip before profiling")
    batches_per_cycle: int = field(10, help="Batches per profiling cycle")
    num_cycles: int = field(1, help="Number of profiling cycles")

    @staticmethod
    def get_type() -> ProfilingScheduleType:
        return ProfilingScheduleType.CYCLICAL


# =============================================================================
# Non-polymorphic nested configs inside MetricsConfig (mirroring real vajra)
# =============================================================================
@frozen_dataclass
class WandbConfig:
    """Weights & Biases integration configuration."""

    enabled: bool = field(False, help="Enable W&B logging")
    project: str = field("vajra", help="W&B project name")


@frozen_dataclass
class ChromeTraceConfig:
    """Chrome tracing configuration."""

    enabled: bool = field(False, help="Enable Chrome tracing output")
    enable_cpu_ops: bool = field(False, help="Include CPU ops in trace")


@frozen_dataclass
class GpuProfilingConfig:
    """GPU profiling configuration."""

    enabled: bool = field(False, help="Enable CUDA profiling")
    output_dir: str = field("./gpu_traces", help="Profiler output directory")
    schedule: AbstractProfilingSchedule = dataclass_field(
        default_factory=BatchwiseProfilingSchedule
    )


@frozen_dataclass
class MetricsConfig:
    """Metrics configuration (non-polymorphic with deeply nested poly)."""

    write_metrics: bool = field(False, help="Write metrics to disk")
    wandb: WandbConfig = dataclass_field(default_factory=WandbConfig)
    chrome_trace: ChromeTraceConfig = dataclass_field(default_factory=ChromeTraceConfig)
    gpu_profiling: GpuProfilingConfig = dataclass_field(
        default_factory=GpuProfilingConfig
    )
    enable_cpu_op_level: bool = field(False, help="Collect CPU op-level metrics")
    enable_per_batch_metrics: bool = field(False, help="Retain per-batch metrics")
    enable_png_output: bool = field(False, help="Save plots as PNG files")


@frozen_dataclass
class InferenceEngineConfig:
    """Top-level inference engine configuration."""

    controller_config: AbstractReplicasetControllerConfig = dataclass_field(
        default_factory=_get_default_replicaset_controller
    )
    metrics_config: MetricsConfig = dataclass_field(default_factory=MetricsConfig)
    resource_allocator_config: AbstractResourceAllocatorConfig = dataclass_field(
        default_factory=LocalGpuResourceAllocatorConfig
    )


# =============================================================================
# Additional patterns from vajra: Computed fields, file paths, validation
# =============================================================================
@frozen_dataclass
class ParallelismConfig:
    """Config with computed fields and divisibility validation (vajra pattern)."""

    tensor_parallel: int = field(1, help="Tensor parallel degree")
    pipeline_parallel: int = field(1, help="Pipeline parallel degree")
    data_parallel: int = field(1, help="Data parallel degree")
    # Computed field set in __post_init__
    _total_gpus: int = dataclass_field(default=0, repr=False)

    def __post_init__(self):
        # Positive validation
        if self.tensor_parallel <= 0:
            raise ValueError("tensor_parallel must be positive")
        if self.pipeline_parallel <= 0:
            raise ValueError("pipeline_parallel must be positive")
        if self.data_parallel <= 0:
            raise ValueError("data_parallel must be positive")
        # Compute derived field (vajra pattern)
        object.__setattr__(
            self,
            "_total_gpus",
            self.tensor_parallel * self.pipeline_parallel * self.data_parallel,
        )

    @property
    def total_gpus(self) -> int:
        """Total GPUs required."""
        return self._total_gpus

    @property
    def world_size(self) -> int:
        """Alias for total_gpus."""
        return self._total_gpus

    def verify_parallelism(self, available_gpus: int) -> bool:
        """Verify parallelism config fits available resources."""
        return self._total_gpus <= available_gpus


@frozen_dataclass
class BatchConfig:
    """Config with cross-field validation and divisibility checks (vajra pattern)."""

    min_batch_size: int = field(1, help="Minimum batch size")
    max_batch_size: int = field(64, help="Maximum batch size")
    batch_size_step: int = field(1, help="Batch size increment step")

    def __post_init__(self):
        # Positive validation
        if self.min_batch_size <= 0:
            raise ValueError("min_batch_size must be positive")
        if self.max_batch_size <= 0:
            raise ValueError("max_batch_size must be positive")
        if self.batch_size_step <= 0:
            raise ValueError("batch_size_step must be positive")
        # Cross-field: min < max
        if self.min_batch_size > self.max_batch_size:
            raise ValueError("min_batch_size must be <= max_batch_size")
        # Divisibility check (vajra pattern)
        range_size = self.max_batch_size - self.min_batch_size
        if range_size > 0 and range_size % self.batch_size_step != 0:
            raise ValueError(
                f"batch size range ({range_size}) must be divisible by step ({self.batch_size_step})"
            )


@frozen_dataclass
class FilePathConfig:
    """Config with file path templates (vajra pattern)."""

    base_dir: str = field("/data", help="Base directory")
    model_name: str = field("llama", help="Model name")
    checkpoint_subdir: str = field("checkpoints", help="Checkpoint subdirectory")
    log_subdir: str = field("logs", help="Log subdirectory")

    @property
    def checkpoint_dir(self) -> str:
        """Computed checkpoint directory path."""
        return f"{self.base_dir}/{self.model_name}/{self.checkpoint_subdir}"

    @property
    def log_dir(self) -> str:
        """Computed log directory path."""
        return f"{self.base_dir}/{self.model_name}/{self.log_subdir}"

    def get_checkpoint_path(self, step: int) -> str:
        """Get path for specific checkpoint step."""
        return f"{self.checkpoint_dir}/step_{step}.pt"

    def get_log_path(self, name: str) -> str:
        """Get path for specific log file."""
        return f"{self.log_dir}/{name}.log"


@frozen_dataclass
class MemoryConfig:
    """Config with conditional field modification in __post_init__ (vajra pattern)."""

    gpu_memory_fraction: float = field(0.9, help="Fraction of GPU memory to use")
    cpu_memory_gb: Optional[int] = field(None, help="CPU memory limit (None=auto)")
    swap_space_gb: int = field(0, help="Swap space in GB")
    # Computed field
    _effective_cpu_memory_gb: int = dataclass_field(default=0, repr=False)

    def __post_init__(self):
        if not 0 < self.gpu_memory_fraction <= 1:
            raise ValueError("gpu_memory_fraction must be in (0, 1]")
        if self.swap_space_gb < 0:
            raise ValueError("swap_space_gb must be non-negative")
        # Conditional field modification (vajra pattern)
        # If cpu_memory_gb is None, auto-detect (simulated as 64GB)
        effective = self.cpu_memory_gb if self.cpu_memory_gb is not None else 64
        object.__setattr__(self, "_effective_cpu_memory_gb", effective)

    @property
    def effective_cpu_memory_gb(self) -> int:
        return self._effective_cpu_memory_gb

    @property
    def total_memory_gb(self) -> int:
        """Total memory including swap."""
        return self._effective_cpu_memory_gb + self.swap_space_gb


# =============================================================================
# Polymorphic Hierarchy: Scheduler (with methods using computed fields)
# =============================================================================
@frozen_dataclass
class AbstractSchedulerConfig(BasePolyConfig):
    """Abstract base for scheduler configurations."""

    max_pending_requests: int = field(1000, help="Max pending requests")

    @classmethod
    def get_type(cls) -> SchedulerType:
        raise NotImplementedError()


@frozen_dataclass
class GreedySchedulerConfig(AbstractSchedulerConfig):
    """Greedy scheduler that processes requests immediately."""

    @staticmethod
    def get_type() -> SchedulerType:
        return SchedulerType.GREEDY


@frozen_dataclass
class BalancedSchedulerConfig(AbstractSchedulerConfig):
    """Balanced scheduler with load balancing."""

    balance_interval_ms: int = field(100, help="Load balance interval in ms")
    max_imbalance_ratio: float = field(0.2, help="Max allowed imbalance ratio")

    @staticmethod
    def get_type() -> SchedulerType:
        return SchedulerType.BALANCED

    def __post_init__(self):
        if self.balance_interval_ms <= 0:
            raise ValueError("balance_interval_ms must be positive")
        if not 0 < self.max_imbalance_ratio < 1:
            raise ValueError("max_imbalance_ratio must be in (0, 1)")


@frozen_dataclass
class PrioritySchedulerConfig(AbstractSchedulerConfig):
    """Priority-based scheduler with multiple priority levels."""

    num_priority_levels: int = field(3, help="Number of priority levels")
    preemption_enabled: bool = field(True, help="Enable request preemption")
    starvation_timeout_ms: int = field(5000, help="Starvation prevention timeout")

    @staticmethod
    def get_type() -> SchedulerType:
        return SchedulerType.PRIORITY

    def __post_init__(self):
        if self.num_priority_levels < 1:
            raise ValueError("num_priority_levels must be at least 1")
        if self.starvation_timeout_ms <= 0:
            raise ValueError("starvation_timeout_ms must be positive")


# =============================================================================
# Polymorphic Hierarchy: Storage Backend (file paths and optional fields)
# =============================================================================
@frozen_dataclass
class AbstractStorageBackendConfig(BasePolyConfig):
    """Abstract base for storage backend configurations."""

    @classmethod
    def get_type(cls) -> StorageBackendType:
        raise NotImplementedError()


@frozen_dataclass
class LocalStorageConfig(AbstractStorageBackendConfig):
    """Local filesystem storage."""

    root_path: str = field("/data/storage", help="Root storage path")
    max_size_gb: Optional[int] = field(None, help="Max size limit (None=unlimited)")

    @staticmethod
    def get_type() -> StorageBackendType:
        return StorageBackendType.LOCAL

    def get_full_path(self, relative_path: str) -> str:
        """Get full path from relative path."""
        return f"{self.root_path}/{relative_path}"


@frozen_dataclass
class S3StorageConfig(AbstractStorageBackendConfig):
    """AWS S3 storage backend."""

    bucket: str = field("my-bucket", help="S3 bucket name")
    prefix: str = field("", help="Key prefix")
    region: str = field("us-east-1", help="AWS region")
    endpoint_url: Optional[str] = field(None, help="Custom endpoint URL")

    @staticmethod
    def get_type() -> StorageBackendType:
        return StorageBackendType.S3

    def get_full_key(self, key: str) -> str:
        """Get full S3 key with prefix."""
        if self.prefix:
            return f"{self.prefix}/{key}"
        return key


@frozen_dataclass
class GCSStorageConfig(AbstractStorageBackendConfig):
    """Google Cloud Storage backend."""

    bucket: str = field("my-bucket", help="GCS bucket name")
    prefix: str = field("", help="Object prefix")
    project: Optional[str] = field(None, help="GCP project ID")

    @staticmethod
    def get_type() -> StorageBackendType:
        return StorageBackendType.GCS


# =============================================================================
# Complex nested config with multiple computed fields
# =============================================================================
@frozen_dataclass
class TrainingConfig:
    """Complex training config with multiple computed fields (vajra pattern)."""

    parallelism: ParallelismConfig = dataclass_field(default_factory=ParallelismConfig)
    batch: BatchConfig = dataclass_field(default_factory=BatchConfig)
    memory: MemoryConfig = dataclass_field(default_factory=MemoryConfig)
    paths: FilePathConfig = dataclass_field(default_factory=FilePathConfig)
    scheduler: AbstractSchedulerConfig = dataclass_field(
        default_factory=GreedySchedulerConfig
    )
    storage: AbstractStorageBackendConfig = dataclass_field(
        default_factory=LocalStorageConfig
    )
    # Training hyperparameters
    learning_rate: float = field(1e-4, help="Learning rate")
    num_epochs: int = field(10, help="Number of training epochs")
    gradient_accumulation_steps: int = field(1, help="Gradient accumulation steps")
    # Computed field
    _effective_batch_size: int = dataclass_field(default=0, repr=False)

    def __post_init__(self):
        if self.learning_rate <= 0:
            raise ValueError("learning_rate must be positive")
        if self.num_epochs <= 0:
            raise ValueError("num_epochs must be positive")
        if self.gradient_accumulation_steps <= 0:
            raise ValueError("gradient_accumulation_steps must be positive")
        # Compute effective batch size (vajra pattern)
        effective = (
            self.batch.max_batch_size
            * self.parallelism.data_parallel
            * self.gradient_accumulation_steps
        )
        object.__setattr__(self, "_effective_batch_size", effective)

    @property
    def effective_batch_size(self) -> int:
        """Effective batch size across all data parallel workers."""
        return self._effective_batch_size

    def get_checkpoint_path(self, epoch: int) -> str:
        """Get checkpoint path for specific epoch."""
        return self.paths.get_checkpoint_path(epoch)

    def validate_resources(self, available_gpus: int) -> bool:
        """Validate that training config fits available resources."""
        return self.parallelism.verify_parallelism(available_gpus)


# =============================================================================
# Config with list of polymorphic configs
# =============================================================================
@frozen_dataclass
class MultiModelConfig:
    """Config with list of polymorphic model configs."""

    models: List[AbstractModelDeploymentConfig] = dataclass_field(default_factory=list)
    default_model_index: int = field(0, help="Index of default model")

    def __post_init__(self):
        if self.models and not 0 <= self.default_model_index < len(self.models):
            raise ValueError(
                f"default_model_index {self.default_model_index} out of range"
            )

    @property
    def default_model(self) -> Optional[AbstractModelDeploymentConfig]:
        """Get the default model config."""
        if not self.models:
            return None
        return self.models[self.default_model_index]

    @property
    def total_world_size(self) -> int:
        """Total world size across all models."""
        return sum(m.world_size for m in self.models)


# =============================================================================
# Tests
# =============================================================================
class TestVajraConfigHierarchy:
    """Test the vajra-like config hierarchy."""

    def test_default_instantiation(self):
        """Test that all configs can be instantiated with defaults."""
        config = InferenceEngineConfig()

        # Check nested structure
        assert isinstance(config.controller_config, AbstractReplicasetControllerConfig)
        assert isinstance(config.controller_config, ReplicasetControllerConfig)
        assert isinstance(
            config.controller_config.replica_controller_config,
            ReplicaControllerConfig,
        )
        assert isinstance(
            config.controller_config.replica_controller_config.model_deployment_config,
            ModelDeploymentConfig,
        )

    def test_polymorphic_type_resolution(self):
        """Test polymorphic type resolution via create_from_type."""
        # Allocator
        lru = AbstractAllocatorConfig.create_from_type(AllocatorType.LRU)
        assert isinstance(lru, LruAllocatorConfig)

        lfu = AbstractAllocatorConfig.create_from_type(AllocatorType.LFU)
        assert isinstance(lfu, LfuAllocatorConfig)

        # Compute manager
        fixed = AbstractComputeManagerConfig.create_from_type(
            ComputeManagerType.FIXED_CHUNK
        )
        assert isinstance(fixed, FixedChunkComputeManagerConfig)

        # Resource allocator
        ray = AbstractResourceAllocatorConfig.create_from_type(
            ResourceAllocatorType.RAY
        )
        assert isinstance(ray, RayGpuResourceAllocatorConfig)

    def test_create_from_dict_simple(self):
        """Test creating configs from dict."""
        config_dict = {
            "type": "cost_aware_exp_decay",
            "prefill_profiling_path": "/custom/path.json",
            "decay_rate": 0.8,
        }
        config = create_class_from_dict(AbstractAllocatorConfig, config_dict)
        assert isinstance(config, CostAwareExpDecayAllocatorConfig)
        assert config.prefill_profiling_path == "/custom/path.json"
        assert config.decay_rate == 0.8

    def test_create_from_dict_nested(self):
        """Test creating nested configs from dict."""
        config_dict = {
            "page_size": 32,
            "allocator": {"type": "lfu"},
            "gpu_tier": {"admission_watermark": 0.1, "demotion_watermark": 0.2},
        }
        config = create_class_from_dict(CacheConfig, config_dict)
        assert config.page_size == 32
        assert isinstance(config.allocator, LfuAllocatorConfig)
        assert config.gpu_tier.admission_watermark == 0.1

    def test_create_from_dict_deep_nesting(self):
        """Test creating deeply nested configs from dict."""
        config_dict = {
            "controller_config": {
                "type": "llm",
                "replica_controller_config": {
                    "type": "llm_base",
                    "model_deployment_config": {
                        "type": "llm",
                        "model": "custom-model",
                        "tensor_parallel_size": 4,
                    },
                    "cache_config": {
                        "page_size": 64,
                        "allocator": {"type": "cost_aware_gdsf"},
                    },
                },
            },
            "resource_allocator_config": {"type": "mock", "num_mock_gpus": 8},
        }

        config = create_class_from_dict(InferenceEngineConfig, config_dict)

        # Verify deep nesting
        assert (
            config.controller_config.replica_controller_config.model_deployment_config.model
            == "custom-model"
        )
        assert (
            config.controller_config.replica_controller_config.model_deployment_config.tensor_parallel_size
            == 4
        )
        assert (
            config.controller_config.replica_controller_config.cache_config.page_size
            == 64
        )
        assert isinstance(
            config.controller_config.replica_controller_config.cache_config.allocator,
            CostAwareGdsfAllocatorConfig,
        )
        assert isinstance(config.resource_allocator_config, MockResourceAllocatorConfig)
        assert config.resource_allocator_config.num_mock_gpus == 8

    def test_dataclass_to_dict_roundtrip(self):
        """Test conversion to dict and back."""
        original = InferenceEngineConfig()
        as_dict = dataclass_to_dict(original)

        # Verify structure
        assert "controller_config" in as_dict
        assert "type" in as_dict["controller_config"]

        # Roundtrip
        reconstructed = create_class_from_dict(InferenceEngineConfig, as_dict)
        assert (
            reconstructed.controller_config.replica_controller_config.model_deployment_config.model
            == original.controller_config.replica_controller_config.model_deployment_config.model
        )

    def test_validation_in_post_init(self):
        """Test that validation in __post_init__ works."""
        # Valid
        config = CostAwareExpDecayAllocatorConfig(decay_rate=0.5)
        assert config.decay_rate == 0.5

        # Invalid
        with pytest.raises(ValueError, match="decay_rate must be positive"):
            CostAwareExpDecayAllocatorConfig(decay_rate=-0.5)

        with pytest.raises(ValueError, match="admission_watermark"):
            CacheTierConfig(admission_watermark=1.5)

    def test_computed_properties(self):
        """Test that computed properties work."""
        config = ModelDeploymentConfig(tensor_parallel_size=4, pipeline_parallel_size=2)
        assert config.world_size == 8

    def test_immutability(self):
        """Test that configs are frozen."""
        config = InferenceEngineConfig()
        with pytest.raises(AttributeError):
            config.metrics_config = MetricsConfig()


class TestVajraConfigCLI:
    """Test CLI argument generation for vajra-like configs."""

    def test_walker_creation(self):
        """Test that ConfigWalker can be created for InferenceEngineConfig."""
        walker = ConfigWalker(InferenceEngineConfig)
        parser = walker.build_parser()
        assert parser is not None

    def test_cli_args_parsing(self):
        """Test that CLI args are generated for nested config."""
        walker = ConfigWalker(InferenceEngineConfig)
        parser = walker.build_parser()
        arg_names = set(parser.arguments.keys())

        # Should have nested config references for polymorphic fields
        assert any("controller_config" in name for name in arg_names)
        assert any("resource_allocator_config" in name for name in arg_names)

    def test_polymorphic_type_selection_cli(self):
        """Test polymorphic type selection via CLI."""
        walker = ConfigWalker(CacheConfig)
        parser = walker.build_parser()
        arg_names = set(parser.arguments.keys())

        # The parser should have a type field for allocator
        assert "allocator.type" in arg_names


class TestVajraConfigEdgeCases:
    """Test edge cases in the config hierarchy."""

    def test_all_allocator_types(self):
        """Test all allocator type variants."""
        for alloc_type in AllocatorType:
            config = AbstractAllocatorConfig.create_from_type(alloc_type)
            assert config.get_type() == alloc_type

    def test_all_compute_manager_types(self):
        """Test all compute manager variants."""
        for cm_type in ComputeManagerType:
            config = AbstractComputeManagerConfig.create_from_type(cm_type)
            assert config.get_type() == cm_type

    def test_string_type_resolution(self):
        """Test that string type values work."""
        config = AbstractAllocatorConfig.create_from_type("lru")
        assert isinstance(config, LruAllocatorConfig)

        config = AbstractAllocatorConfig.create_from_type("LRU")  # case insensitive
        assert isinstance(config, LruAllocatorConfig)

    def test_empty_polymorphic_config(self):
        """Test polymorphic configs with no extra fields."""
        config = LruAllocatorConfig()
        assert config.get_type() == AllocatorType.LRU

        config = PullSessionRouterConfig()
        assert config.get_type() == SessionRouterType.PULL

    def test_nested_default_factories(self):
        """Test that nested default factories work correctly."""
        config = InferenceEngineConfig()

        # Each default should be a fresh instance
        config2 = InferenceEngineConfig()

        # They should be equal but not the same object
        assert (
            config.controller_config.replica_controller_config.model_deployment_config.model
            == config2.controller_config.replica_controller_config.model_deployment_config.model
        )


class TestComputedFields:
    """Test computed fields set in __post_init__ (vajra pattern)."""

    def test_parallelism_computed_total(self):
        """Test computed total_gpus field."""
        config = ParallelismConfig(
            tensor_parallel=4, pipeline_parallel=2, data_parallel=8
        )
        assert config.total_gpus == 64
        assert config.world_size == 64

    def test_parallelism_verify_method(self):
        """Test method using computed fields."""
        config = ParallelismConfig(tensor_parallel=4, pipeline_parallel=2)
        assert config.verify_parallelism(available_gpus=8)
        assert not config.verify_parallelism(available_gpus=4)

    def test_memory_conditional_field(self):
        """Test conditional field modification in __post_init__."""
        # Auto-detect (None -> 64)
        config = MemoryConfig(cpu_memory_gb=None)
        assert config.effective_cpu_memory_gb == 64
        assert config.total_memory_gb == 64

        # Explicit value
        config = MemoryConfig(cpu_memory_gb=32, swap_space_gb=16)
        assert config.effective_cpu_memory_gb == 32
        assert config.total_memory_gb == 48

    def test_training_effective_batch_size(self):
        """Test effective batch size computed from nested configs."""
        config = TrainingConfig(
            batch=BatchConfig(max_batch_size=32),
            parallelism=ParallelismConfig(data_parallel=4),
            gradient_accumulation_steps=2,
        )
        # 32 * 4 * 2 = 256
        assert config.effective_batch_size == 256

    def test_multi_model_total_world_size(self):
        """Test computed property across list of configs."""
        config = MultiModelConfig(
            models=[
                ModelDeploymentConfig(tensor_parallel_size=4),
                ModelDeploymentConfig(tensor_parallel_size=2, pipeline_parallel_size=2),
            ]
        )
        # 4*1 + 2*2 = 8
        assert config.total_world_size == 8
        assert config.default_model.tensor_parallel_size == 4


class TestCrossFieldValidation:
    """Test cross-field validation (vajra pattern)."""

    def test_dynamic_chunk_min_max_validation(self):
        """Test min < max validation."""
        # Valid
        config = DynamicChunkComputeManagerConfig(
            min_chunk_size=512, max_chunk_size=4096
        )
        assert config.min_chunk_size < config.max_chunk_size

        # Invalid: min >= max
        with pytest.raises(ValueError, match="min_chunk_size must be less than"):
            DynamicChunkComputeManagerConfig(min_chunk_size=4096, max_chunk_size=512)

        with pytest.raises(ValueError, match="min_chunk_size must be less than"):
            DynamicChunkComputeManagerConfig(min_chunk_size=1000, max_chunk_size=1000)

    def test_cache_tier_watermark_validation(self):
        """Test admission < demotion validation."""
        # Valid
        config = CacheTierConfig(admission_watermark=0.05, demotion_watermark=0.1)
        assert config.admission_watermark < config.demotion_watermark

        # Invalid: admission >= demotion
        with pytest.raises(ValueError, match="admission_watermark must be less than"):
            CacheTierConfig(admission_watermark=0.2, demotion_watermark=0.1)

    def test_batch_divisibility_validation(self):
        """Test divisibility validation."""
        # Valid: range 63 is divisible by step 1
        config = BatchConfig(min_batch_size=1, max_batch_size=64, batch_size_step=1)
        assert config.min_batch_size == 1

        # Valid: range 56 is divisible by step 8
        config = BatchConfig(min_batch_size=8, max_batch_size=64, batch_size_step=8)
        assert config.batch_size_step == 8

        # Invalid: range 60 is not divisible by step 8
        with pytest.raises(ValueError, match="divisible by step"):
            BatchConfig(min_batch_size=4, max_batch_size=64, batch_size_step=8)

    def test_batch_min_max_validation(self):
        """Test min <= max validation."""
        with pytest.raises(ValueError, match="min_batch_size must be <="):
            BatchConfig(min_batch_size=100, max_batch_size=50)


class TestFilePathTemplates:
    """Test file path template patterns (vajra pattern)."""

    def test_computed_paths(self):
        """Test computed directory paths."""
        config = FilePathConfig(
            base_dir="/models", model_name="llama-7b", checkpoint_subdir="ckpt"
        )
        assert config.checkpoint_dir == "/models/llama-7b/ckpt"
        assert config.log_dir == "/models/llama-7b/logs"

    def test_path_methods(self):
        """Test path generation methods."""
        config = FilePathConfig(base_dir="/data", model_name="gpt")
        assert config.get_checkpoint_path(100) == "/data/gpt/checkpoints/step_100.pt"
        assert config.get_log_path("train") == "/data/gpt/logs/train.log"

    def test_training_config_path_delegation(self):
        """Test path method delegation in nested config."""
        config = TrainingConfig(
            paths=FilePathConfig(base_dir="/experiments", model_name="test")
        )
        assert (
            config.get_checkpoint_path(50) == "/experiments/test/checkpoints/step_50.pt"
        )


class TestSchedulerHierarchy:
    """Test scheduler polymorphic hierarchy."""

    def test_all_scheduler_types(self):
        """Test all scheduler type variants."""
        for sched_type in SchedulerType:
            config = AbstractSchedulerConfig.create_from_type(sched_type)
            assert config.get_type() == sched_type

    def test_scheduler_from_dict(self):
        """Test creating schedulers from dict."""
        config = create_class_from_dict(
            AbstractSchedulerConfig,
            {
                "type": "balanced",
                "balance_interval_ms": 200,
                "max_imbalance_ratio": 0.1,
            },
        )
        assert isinstance(config, BalancedSchedulerConfig)
        assert config.balance_interval_ms == 200

    def test_scheduler_validation(self):
        """Test scheduler-specific validation."""
        with pytest.raises(ValueError, match="balance_interval_ms must be positive"):
            BalancedSchedulerConfig(balance_interval_ms=0)

        with pytest.raises(ValueError, match="num_priority_levels must be at least 1"):
            PrioritySchedulerConfig(num_priority_levels=0)


class TestStorageBackendHierarchy:
    """Test storage backend polymorphic hierarchy."""

    def test_all_storage_types(self):
        """Test all storage backend type variants."""
        for storage_type in StorageBackendType:
            config = AbstractStorageBackendConfig.create_from_type(storage_type)
            assert config.get_type() == storage_type

    def test_local_storage_path_method(self):
        """Test local storage path method."""
        config = LocalStorageConfig(root_path="/data/store")
        assert config.get_full_path("models/v1") == "/data/store/models/v1"

    def test_s3_storage_key_method(self):
        """Test S3 storage key method."""
        config = S3StorageConfig(bucket="my-bucket", prefix="training/run1")
        assert config.get_full_key("checkpoint.pt") == "training/run1/checkpoint.pt"

        config_no_prefix = S3StorageConfig(bucket="my-bucket", prefix="")
        assert config_no_prefix.get_full_key("checkpoint.pt") == "checkpoint.pt"

    def test_storage_from_dict(self):
        """Test creating storage configs from dict."""
        config = create_class_from_dict(
            AbstractStorageBackendConfig,
            {"type": "s3", "bucket": "test-bucket", "region": "eu-west-1"},
        )
        assert isinstance(config, S3StorageConfig)
        assert config.bucket == "test-bucket"
        assert config.region == "eu-west-1"


class TestTrainingConfigIntegration:
    """Integration tests for complex TrainingConfig."""

    def test_default_instantiation(self):
        """Test TrainingConfig with all defaults."""
        config = TrainingConfig()
        assert config.parallelism.total_gpus == 1
        assert config.effective_batch_size == 64  # 64 * 1 * 1
        assert isinstance(config.scheduler, GreedySchedulerConfig)
        assert isinstance(config.storage, LocalStorageConfig)

    def test_full_config_from_dict(self):
        """Test creating full TrainingConfig from dict."""
        config_dict = {
            "parallelism": {
                "tensor_parallel": 4,
                "pipeline_parallel": 2,
                "data_parallel": 2,
            },
            "batch": {"min_batch_size": 8, "max_batch_size": 32, "batch_size_step": 8},
            "memory": {"gpu_memory_fraction": 0.95, "cpu_memory_gb": 128},
            "paths": {"base_dir": "/experiments", "model_name": "my-model"},
            "scheduler": {
                "type": "priority",
                "num_priority_levels": 5,
                "preemption_enabled": True,
            },
            "storage": {"type": "s3", "bucket": "training-data", "prefix": "run123"},
            "learning_rate": 0.0001,
            "num_epochs": 100,
            "gradient_accumulation_steps": 4,
        }
        config = create_class_from_dict(TrainingConfig, config_dict)

        # Verify nested configs
        assert config.parallelism.total_gpus == 16  # 4*2*2
        assert config.parallelism.verify_parallelism(16)
        assert not config.parallelism.verify_parallelism(8)

        # Verify computed fields
        # effective_batch = 32 * 2 * 4 = 256
        assert config.effective_batch_size == 256

        # Verify polymorphic configs
        assert isinstance(config.scheduler, PrioritySchedulerConfig)
        assert config.scheduler.num_priority_levels == 5
        assert isinstance(config.storage, S3StorageConfig)
        assert config.storage.bucket == "training-data"

        # Verify path methods
        assert config.paths.checkpoint_dir == "/experiments/my-model/checkpoints"

    def test_roundtrip_serialization(self):
        """Test TrainingConfig roundtrip through dict."""
        original = TrainingConfig(
            parallelism=ParallelismConfig(tensor_parallel=8),
            scheduler=PrioritySchedulerConfig(num_priority_levels=4),
            storage=S3StorageConfig(bucket="test"),
            learning_rate=0.001,
        )
        as_dict = dataclass_to_dict(original)
        reconstructed = create_class_from_dict(TrainingConfig, as_dict)

        assert reconstructed.parallelism.tensor_parallel == 8
        assert reconstructed.parallelism.total_gpus == 8
        assert isinstance(reconstructed.scheduler, PrioritySchedulerConfig)
        assert reconstructed.learning_rate == 0.001

    def test_validate_resources_method(self):
        """Test resource validation method."""
        config = TrainingConfig(
            parallelism=ParallelismConfig(tensor_parallel=4, data_parallel=2)
        )
        assert config.validate_resources(available_gpus=8)
        assert not config.validate_resources(available_gpus=4)


class TestMultiModelConfig:
    """Test config with list of polymorphic configs."""

    def test_empty_models_list(self):
        """Test MultiModelConfig with empty models."""
        config = MultiModelConfig(models=[])
        assert config.default_model is None
        assert config.total_world_size == 0

    def test_models_list_from_dict(self):
        """Test creating MultiModelConfig with models from dict."""
        config_dict = {
            "models": [
                {"type": "llm", "model": "llama-7b", "tensor_parallel_size": 2},
                {"type": "llm", "model": "llama-13b", "tensor_parallel_size": 4},
            ],
            "default_model_index": 1,
        }
        config = create_class_from_dict(MultiModelConfig, config_dict)

        assert len(config.models) == 2
        assert config.models[0].model == "llama-7b"
        assert config.models[1].model == "llama-13b"
        assert config.default_model.model == "llama-13b"
        assert config.total_world_size == 6  # 2 + 4

    def test_invalid_default_index(self):
        """Test validation of default_model_index."""
        with pytest.raises(ValueError, match="default_model_index .* out of range"):
            MultiModelConfig(
                models=[ModelDeploymentConfig()],
                default_model_index=5,
            )


class TestPositiveValidation:
    """Test positive value validation across configs."""

    def test_parallelism_positive(self):
        """Test parallelism values must be positive."""
        with pytest.raises(ValueError, match="tensor_parallel must be positive"):
            ParallelismConfig(tensor_parallel=0)

        with pytest.raises(ValueError, match="pipeline_parallel must be positive"):
            ParallelismConfig(pipeline_parallel=-1)

    def test_batch_positive(self):
        """Test batch values must be positive."""
        with pytest.raises(ValueError, match="min_batch_size must be positive"):
            BatchConfig(min_batch_size=0)

        with pytest.raises(ValueError, match="batch_size_step must be positive"):
            BatchConfig(batch_size_step=0)

    def test_memory_validation(self):
        """Test memory config validation."""
        with pytest.raises(ValueError, match="gpu_memory_fraction must be in"):
            MemoryConfig(gpu_memory_fraction=0)

        with pytest.raises(ValueError, match="gpu_memory_fraction must be in"):
            MemoryConfig(gpu_memory_fraction=1.5)

        with pytest.raises(ValueError, match="swap_space_gb must be non-negative"):
            MemoryConfig(swap_space_gb=-1)

    def test_training_positive(self):
        """Test training config positive validation."""
        with pytest.raises(ValueError, match="learning_rate must be positive"):
            TrainingConfig(learning_rate=0)

        with pytest.raises(ValueError, match="num_epochs must be positive"):
            TrainingConfig(num_epochs=0)


class TestWalkerForNewConfigs:
    """Test walker and parser for new config types."""

    def test_training_config_walker(self):
        """Test walker for TrainingConfig."""
        walker = ConfigWalker(TrainingConfig)
        parser = walker.build_parser()
        arg_names = set(parser.arguments.keys())

        # Should have nested fields with prefixes
        assert any("parallelism" in name for name in arg_names)
        assert any("batch" in name for name in arg_names)
        assert any("scheduler" in name for name in arg_names)

    def test_multi_model_walker(self):
        """Test walker for MultiModelConfig."""
        walker = ConfigWalker(MultiModelConfig)
        parser = walker.build_parser()
        assert parser is not None


# =============================================================================
# CLI Argument Parsing Tests (vajra pattern)
# =============================================================================


class TestCLIArgumentParsing:
    """Test CLI argument parsing for configs."""

    def test_parser_arg_names(self):
        """Test that parser has correct argument names."""
        walker = ConfigWalker(CacheConfig)
        parser = walker.build_parser()
        arg_names = set(parser.arguments.keys())

        # Polymorphic type field
        assert "allocator.type" in arg_names
        # Nested tier fields
        assert "gpu_tier.admission_watermark" in arg_names
        assert "gpu_tier.demotion_watermark" in arg_names

    def test_deep_nesting_parser(self):
        """Test parser for deeply nested configs."""
        walker = ConfigWalker(InferenceEngineConfig)
        parser = walker.build_parser()
        arg_names = set(parser.arguments.keys())

        # Should have nested polymorphic config fields
        assert any("controller_config" in name for name in arg_names)
        assert any("replica_controller_config" in name for name in arg_names)

    def test_polymorphic_type_field(self):
        """Test that polymorphic configs get type fields."""
        walker = ConfigWalker(CacheConfig)
        parser = walker.build_parser()
        arg_names = set(parser.arguments.keys())

        # Polymorphic allocator should have a type field
        assert "allocator.type" in arg_names

    def test_parser_has_all_nested_fields(self):
        """Test that parser has all fields from nested structure."""
        walker = ConfigWalker(CacheConfig)
        parser = walker.build_parser()
        arg_names = set(parser.arguments.keys())

        # Should have top-level fields
        assert "page_size" in arg_names
        assert "max_promotion_batch_size" in arg_names
        assert "enable_prefix_caching" in arg_names

        # Should have nested tier fields
        assert "gpu_tier.admission_watermark" in arg_names
        assert "gpu_tier.demotion_watermark" in arg_names
        assert "cpu_tier.admission_watermark" in arg_names
        assert "nvme_tier.demotion_watermark" in arg_names

        # Should have polymorphic type field
        assert "allocator.type" in arg_names

        # Should have variant-specific fields
        assert any("decay_rate" in name for name in arg_names)
        assert any("use_frequency" in name for name in arg_names)

    def test_cli_help_text_propagation(self):
        """Test that help text from field() is preserved in parser."""
        walker = ConfigWalker(CacheConfig)
        parser = walker.build_parser()

        # Check that help text is preserved
        page_size_arg = parser.arguments.get("page_size")
        assert page_size_arg is not None
        assert page_size_arg.help_text is not None
        assert "Cache page size" in page_size_arg.help_text

    def test_training_config_parser(self):
        """Test parser for complex TrainingConfig."""
        walker = ConfigWalker(TrainingConfig)
        parser = walker.build_parser()
        arg_names = set(parser.arguments.keys())

        # Check various nested fields exist
        assert "learning_rate" in arg_names
        assert "num_epochs" in arg_names
        assert "parallelism.tensor_parallel" in arg_names
        assert "parallelism.pipeline_parallel" in arg_names
        assert "batch.min_batch_size" in arg_names
        assert "batch.max_batch_size" in arg_names
        assert "scheduler.type" in arg_names
        assert "storage.type" in arg_names


class TestFileBasedConfigLoading:
    """Test YAML/JSON file-based config loading (vajra pattern)."""

    def test_load_yaml_simple(self):
        """Test loading simple YAML config."""
        yaml_content = """
page_size: 64
enable_prefix_caching: false
allocator:
  type: lfu
"""
        with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f:
            f.write(yaml_content)
            yaml_path = f.name

        try:
            config_dict = load_yaml_config(yaml_path)
            config = create_class_from_dict(CacheConfig, config_dict)

            assert config.page_size == 64
            assert config.enable_prefix_caching is False
            assert isinstance(config.allocator, LfuAllocatorConfig)
        finally:
            os.unlink(yaml_path)

    def test_load_yaml_nested(self):
        """Test loading nested YAML config."""
        yaml_content = """
controller_config:
  type: llm
  replica_controller_config:
    type: llm_base
    model_deployment_config:
      type: llm
      model: "gpt-4"
      tensor_parallel_size: 8
    cache_config:
      page_size: 128
      allocator:
        type: cost_aware_exp_decay
        decay_rate: 0.7
metrics_config:
  write_metrics: true
  enable_cpu_op_level: true
  wandb:
    enabled: true
    project: "my_project"
  gpu_profiling:
    enabled: true
    schedule:
      type: cyclical
      warmup_batches: 5
      batches_per_cycle: 20
resource_allocator_config:
  type: ray
  ray_address: "ray://cluster:10001"
"""
        with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f:
            f.write(yaml_content)
            yaml_path = f.name

        try:
            config_dict = load_yaml_config(yaml_path)
            config = create_class_from_dict(InferenceEngineConfig, config_dict)

            # Verify deep nesting
            assert (
                config.controller_config.replica_controller_config.model_deployment_config.model
                == "gpt-4"
            )
            assert (
                config.controller_config.replica_controller_config.model_deployment_config.tensor_parallel_size
                == 8
            )
            assert (
                config.controller_config.replica_controller_config.cache_config.page_size
                == 128
            )
            assert isinstance(
                config.controller_config.replica_controller_config.cache_config.allocator,
                CostAwareExpDecayAllocatorConfig,
            )
            # Verify metrics with deeply nested poly
            assert config.metrics_config.write_metrics is True
            assert config.metrics_config.enable_cpu_op_level is True
            assert config.metrics_config.wandb.enabled is True
            assert config.metrics_config.wandb.project == "my_project"
            assert config.metrics_config.gpu_profiling.enabled is True
            assert isinstance(
                config.metrics_config.gpu_profiling.schedule,
                CyclicalProfilingSchedule,
            )
            assert config.metrics_config.gpu_profiling.schedule.warmup_batches == 5
            assert config.metrics_config.gpu_profiling.schedule.batches_per_cycle == 20
            assert isinstance(
                config.resource_allocator_config, RayGpuResourceAllocatorConfig
            )
        finally:
            os.unlink(yaml_path)

    def test_load_yaml_with_training_config(self):
        """Test loading complex TrainingConfig from YAML."""
        yaml_content = """
parallelism:
  tensor_parallel: 4
  pipeline_parallel: 2
  data_parallel: 2
batch:
  min_batch_size: 8
  max_batch_size: 64
  batch_size_step: 8
memory:
  gpu_memory_fraction: 0.95
  cpu_memory_gb: 256
paths:
  base_dir: "/experiments"
  model_name: "llama-70b"
scheduler:
  type: priority
  num_priority_levels: 5
storage:
  type: s3
  bucket: "ml-training"
  prefix: "checkpoints"
learning_rate: 0.0001
num_epochs: 50
gradient_accumulation_steps: 8
"""
        with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f:
            f.write(yaml_content)
            yaml_path = f.name

        try:
            config_dict = load_yaml_config(yaml_path)
            config = create_class_from_dict(TrainingConfig, config_dict)

            # Verify all nested values
            assert config.parallelism.tensor_parallel == 4
            assert config.parallelism.total_gpus == 16  # 4*2*2
            assert config.batch.max_batch_size == 64
            assert config.memory.effective_cpu_memory_gb == 256
            assert config.paths.checkpoint_dir == "/experiments/llama-70b/checkpoints"
            assert isinstance(config.scheduler, PrioritySchedulerConfig)
            assert config.scheduler.num_priority_levels == 5
            assert isinstance(config.storage, S3StorageConfig)
            assert config.storage.bucket == "ml-training"
            # Verify computed field
            # effective_batch = 64 * 2 * 8 = 1024
            assert config.effective_batch_size == 1024
        finally:
            os.unlink(yaml_path)

    def test_load_json_config(self):
        """Test loading JSON config file."""
        json_content = '{"page_size": 32, "allocator": {"type": "lru"}}'

        with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f:
            f.write(json_content)
            json_path = f.name

        try:
            config_dict = load_yaml_config(json_path)  # YAML loader handles JSON too
            config = create_class_from_dict(CacheConfig, config_dict)

            assert config.page_size == 32
            assert isinstance(config.allocator, LruAllocatorConfig)
        finally:
            os.unlink(json_path)


class TestConfigRoundtrip:
    """Test config serialization roundtrip (vajra pattern: save -> load -> verify)."""

    def test_simple_config_roundtrip(self):
        """Test simple config YAML roundtrip."""
        original = CacheConfig(
            page_size=64,
            allocator=CostAwareGdsfAllocatorConfig(
                prefill_profiling_path="/custom/path.json"
            ),
            enable_prefix_caching=False,
        )

        # Serialize to dict
        as_dict = dataclass_to_dict(original)

        # Save to YAML
        with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f:
            yaml.dump(as_dict, f)
            yaml_path = f.name

        try:
            # Load back
            loaded_dict = load_yaml_config(yaml_path)
            reconstructed = create_class_from_dict(CacheConfig, loaded_dict)

            # Verify
            assert reconstructed.page_size == original.page_size
            assert reconstructed.enable_prefix_caching == original.enable_prefix_caching
            assert isinstance(reconstructed.allocator, CostAwareGdsfAllocatorConfig)
            assert (
                reconstructed.allocator.prefill_profiling_path
                == original.allocator.prefill_profiling_path
            )
        finally:
            os.unlink(yaml_path)

    def test_complex_config_roundtrip(self):
        """Test complex nested config YAML roundtrip."""
        original = InferenceEngineConfig(
            controller_config=ReplicasetControllerConfig(
                replica_controller_config=ReplicaControllerConfig(
                    model_deployment_config=ModelDeploymentConfig(
                        model="test-model",
                        tensor_parallel_size=4,
                    ),
                    cache_config=CacheConfig(
                        page_size=64,
                        allocator=CostAwareExpDecayAllocatorConfig(decay_rate=0.8),
                    ),
                ),
                session_prioritizer_config=EdfSessionPrioritizerConfig(
                    default_deadline_ms=500.0
                ),
            ),
            resource_allocator_config=MockResourceAllocatorConfig(num_mock_gpus=16),
        )

        # Serialize
        as_dict = dataclass_to_dict(original)

        # Save and load
        with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f:
            yaml.dump(as_dict, f)
            yaml_path = f.name

        try:
            loaded_dict = load_yaml_config(yaml_path)
            reconstructed = create_class_from_dict(InferenceEngineConfig, loaded_dict)

            # Deep verification
            orig_model = (
                original.controller_config.replica_controller_config.model_deployment_config
            )
            recon_model = (
                reconstructed.controller_config.replica_controller_config.model_deployment_config
            )
            assert recon_model.model == orig_model.model
            assert recon_model.tensor_parallel_size == orig_model.tensor_parallel_size

            orig_cache = (
                original.controller_config.replica_controller_config.cache_config
            )
            recon_cache = (
                reconstructed.controller_config.replica_controller_config.cache_config
            )
            assert recon_cache.page_size == orig_cache.page_size
            assert isinstance(recon_cache.allocator, CostAwareExpDecayAllocatorConfig)
            assert recon_cache.allocator.decay_rate == 0.8

            assert isinstance(
                reconstructed.controller_config.session_prioritizer_config,
                EdfSessionPrioritizerConfig,
            )
            assert isinstance(
                reconstructed.resource_allocator_config, MockResourceAllocatorConfig
            )
        finally:
            os.unlink(yaml_path)

    def test_training_config_roundtrip(self):
        """Test TrainingConfig YAML roundtrip with computed fields."""
        original = TrainingConfig(
            parallelism=ParallelismConfig(tensor_parallel=8, data_parallel=4),
            batch=BatchConfig(min_batch_size=16, max_batch_size=64, batch_size_step=16),
            scheduler=BalancedSchedulerConfig(balance_interval_ms=50),
            storage=GCSStorageConfig(bucket="test-bucket", project="my-project"),
            learning_rate=0.0005,
            gradient_accumulation_steps=4,
        )

        # Verify original computed fields
        assert original.parallelism.total_gpus == 32
        assert original.effective_batch_size == 64 * 4 * 4  # 1024

        # Roundtrip
        as_dict = dataclass_to_dict(original)

        with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f:
            yaml.dump(as_dict, f)
            yaml_path = f.name

        try:
            loaded_dict = load_yaml_config(yaml_path)
            reconstructed = create_class_from_dict(TrainingConfig, loaded_dict)

            # Verify values preserved
            assert reconstructed.parallelism.tensor_parallel == 8
            assert reconstructed.parallelism.data_parallel == 4
            assert reconstructed.learning_rate == 0.0005

            # Verify computed fields recomputed correctly
            assert reconstructed.parallelism.total_gpus == 32
            assert reconstructed.effective_batch_size == 1024

            # Verify polymorphic types
            assert isinstance(reconstructed.scheduler, BalancedSchedulerConfig)
            assert reconstructed.scheduler.balance_interval_ms == 50
            assert isinstance(reconstructed.storage, GCSStorageConfig)
            assert reconstructed.storage.project == "my-project"
        finally:
            os.unlink(yaml_path)


class TestConfigFileLoading:
    """Test --config flag for loading configs from files."""

    def test_config_file_loading(self):
        """Test that --config flag loads config from YAML file."""

        @frozen_dataclass
        class FileLoadableConfig:
            value: int = field(10, help="A value")
            name: str = field("default", help="A name")

        yaml_content = """
value: 42
name: "from_file"
"""
        with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f:
            f.write(yaml_content)
            yaml_path = f.name

        try:
            test_args = ["script.py", "--config", yaml_path]

            with patch.object(sys, "argv", test_args):
                config = parse_cli_args(FileLoadableConfig)

            assert config.value == 42
            assert config.name == "from_file"
        finally:
            Path(yaml_path).unlink()

    def test_config_file_with_nested_configs(self):
        """Test --config with nested configs."""

        @frozen_dataclass
        class InnerConfig:
            inner_val: int = 5

        @frozen_dataclass
        class OuterConfig:
            inner: InnerConfig = dataclass_field(default_factory=InnerConfig)
            outer_val: str = "outer"

        yaml_content = """
outer_val: "from_config"
inner:
  inner_val: 99
"""
        with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f:
            f.write(yaml_content)
            yaml_path = f.name

        try:
            test_args = ["script.py", "--config", yaml_path]

            with patch.object(sys, "argv", test_args):
                config = parse_cli_args(OuterConfig)

            assert config.outer_val == "from_config"
            assert config.inner.inner_val == 99
        finally:
            Path(yaml_path).unlink()


class TestPolymorphicCLI:
    """Test polymorphic config handling in CLI."""

    def test_polymorphic_type_field_in_parser(self):
        """Test that polymorphic configs register a .type selector in the parser."""

        @frozen_dataclass
        class ConfigWithPoly:
            allocator: AbstractAllocatorConfig = dataclass_field(
                default_factory=LruAllocatorConfig
            )
            name: str = "test"

        walker = ConfigWalker(ConfigWithPoly)
        parser = walker.build_parser()
        arg_names = set(parser.arguments.keys())

        # Type selector uses .type suffix
        assert "allocator.type" in arg_names

    def test_polymorphic_variant_fields_tracked(self):
        """Test that polymorphic variant fields are tracked separately."""

        @frozen_dataclass
        class ConfigWithPoly:
            allocator: AbstractAllocatorConfig = dataclass_field(
                default_factory=LruAllocatorConfig
            )

        walker = ConfigWalker(ConfigWithPoly)
        walker.build_parser()

        # Variant-specific fields should be tracked
        assert len(walker.poly_variant_fields) > 0
        # Fields like decay_rate (only for cost_aware_exp_decay) should be variant fields
        assert any(
            "decay_rate" in field_name for field_name in walker.poly_variant_fields
        )

    def test_all_polymorphic_variants_have_fields(self):
        """Test that all polymorphic variant fields are present in parser."""

        @frozen_dataclass
        class ConfigWithPoly:
            allocator: AbstractAllocatorConfig = dataclass_field(
                default_factory=LruAllocatorConfig
            )

        walker = ConfigWalker(ConfigWithPoly)
        parser = walker.build_parser()
        arg_names = set(parser.arguments.keys())

        # All allocator variant fields should be registered with . separator
        # CostAwareExpDecayAllocatorConfig fields
        assert any("decay_rate" in name for name in arg_names)
        assert any("min_cost" in name for name in arg_names)
        assert any("prefill_profiling_path" in name for name in arg_names)

        # CostAwareGdsfAllocatorConfig fields
        assert any("use_frequency" in name for name in arg_names)


class TestEnumHandling:
    """Test enum handling in CLI and serialization."""

    def test_enum_serialization(self):
        """Test that enums serialize to their values."""
        # Use allocator config which has enum type
        config = LruAllocatorConfig()
        as_dict = dataclass_to_dict(config)

        # Type should be the enum VALUE, not name
        assert as_dict["type"] == "lru"  # Not "LRU"

    def test_enum_deserialization_from_value(self):
        """Test that enums can be deserialized from string values."""
        config_dict = {"type": "cost_aware_gdsf", "use_frequency": False}
        config = create_class_from_dict(AbstractAllocatorConfig, config_dict)

        assert isinstance(config, CostAwareGdsfAllocatorConfig)
        assert config.use_frequency is False

    def test_enum_deserialization_case_insensitive(self):
        """Test case-insensitive enum deserialization."""
        # Uppercase
        config = AbstractAllocatorConfig.create_from_type("LRU")
        assert isinstance(config, LruAllocatorConfig)

        # Lowercase
        config = AbstractAllocatorConfig.create_from_type("lru")
        assert isinstance(config, LruAllocatorConfig)

        # Mixed case
        config = AbstractAllocatorConfig.create_from_type("Lru")
        assert isinstance(config, LruAllocatorConfig)

    def test_scheduler_enum_roundtrip(self):
        """Test scheduler enum serialization roundtrip."""
        original = TrainingConfig(
            scheduler=PrioritySchedulerConfig(num_priority_levels=7)
        )

        as_dict = dataclass_to_dict(original)

        # Verify enum value in dict
        assert as_dict["scheduler"]["type"] == "priority"

        # Roundtrip
        reconstructed = create_class_from_dict(TrainingConfig, as_dict)
        assert isinstance(reconstructed.scheduler, PrioritySchedulerConfig)
        assert reconstructed.scheduler.num_priority_levels == 7


class TestListFieldHandling:
    """Test list field handling in configs."""

    def test_list_of_primitives(self):
        """Test list of primitive values."""

        @frozen_dataclass
        class ConfigWithList:
            values: List[int] = dataclass_field(default_factory=list)
            name: str = "test"

        config = ConfigWithList(values=[1, 2, 3, 4, 5])
        as_dict = dataclass_to_dict(config)

        assert as_dict["values"] == [1, 2, 3, 4, 5]

        reconstructed = create_class_from_dict(ConfigWithList, as_dict)
        assert reconstructed.values == [1, 2, 3, 4, 5]

    def test_list_of_dataclasses(self):
        """Test list of dataclass values."""
        config = MultiModelConfig(
            models=[
                ModelDeploymentConfig(model="model-a", tensor_parallel_size=2),
                ModelDeploymentConfig(model="model-b", tensor_parallel_size=4),
            ]
        )

        as_dict = dataclass_to_dict(config)

        assert len(as_dict["models"]) == 2
        assert as_dict["models"][0]["model"] == "model-a"

        reconstructed = create_class_from_dict(MultiModelConfig, as_dict)
        assert len(reconstructed.models) == 2
        assert reconstructed.models[0].model == "model-a"
        assert reconstructed.models[1].tensor_parallel_size == 4

    def test_list_field_in_walker(self):
        """Test that list fields are tracked in walker."""
        walker = ConfigWalker(MultiModelConfig)
        walker.build_parser()

        assert len(walker.list_fields) > 0


class TestOptionalFieldHandling:
    """Test Optional field handling."""

    def test_optional_none_value(self):
        """Test Optional field with None value."""
        config = MemoryConfig(cpu_memory_gb=None)

        as_dict = dataclass_to_dict(config)
        # None should be serialized
        assert as_dict["cpu_memory_gb"] is None

        reconstructed = create_class_from_dict(MemoryConfig, as_dict)
        assert reconstructed.cpu_memory_gb is None
        # But effective memory should be auto-detected
        assert reconstructed.effective_cpu_memory_gb == 64

    def test_optional_with_value(self):
        """Test Optional field with actual value."""
        config = MemoryConfig(cpu_memory_gb=128)

        as_dict = dataclass_to_dict(config)
        assert as_dict["cpu_memory_gb"] == 128

        reconstructed = create_class_from_dict(MemoryConfig, as_dict)
        assert reconstructed.cpu_memory_gb == 128
        assert reconstructed.effective_cpu_memory_gb == 128

    def test_optional_in_storage_config(self):
        """Test Optional fields in storage configs."""
        # S3 with optional endpoint
        config = S3StorageConfig(bucket="test", endpoint_url="http://localhost:9000")
        as_dict = dataclass_to_dict(config)
        assert as_dict["endpoint_url"] == "http://localhost:9000"

        # S3 without optional endpoint
        config = S3StorageConfig(bucket="test", endpoint_url=None)
        as_dict = dataclass_to_dict(config)
        assert as_dict["endpoint_url"] is None


# =============================================================================
# Tests for CLI help completeness and non-polymorphic enum handling
# =============================================================================
class TestCLIHelpCompleteness:
    """Regression tests for CLI help showing all config fields.

    These tests guard against bugs where:
    1. Non-polymorphic enum fields (like attention_backend_type) were filtered
       from help because they end with '.type' but aren't type selectors.
    2. Non-polymorphic configs with deeply nested polymorphism (like MetricsConfig
       containing gpu_profiling.schedule) had all their fields dropped from help.
    """

    def test_all_walker_fields_present_in_parser(self):
        """Every field registered by the walker must have a parser argument."""
        walker = ConfigWalker(InferenceEngineConfig)
        parser = walker.build_parser()

        # The parser should have arguments for all walker-discovered fields
        arg_names = set(parser.arguments.keys())
        assert len(arg_names) > 0, "Parser has no arguments"

        # Spot-check key fields are present
        assert "controller_config.type" in arg_names
        assert "resource_allocator_config.type" in arg_names
        assert "metrics_config.write_metrics" in arg_names

    def test_all_fields_present_in_help_output(self):
        """All parser arguments must appear in the rendered help text."""
        walker = ConfigWalker(InferenceEngineConfig)
        parser = walker.build_parser()

        old_stdout = sys.stdout
        sys.stdout = captured = StringIO()
        parser.print_help()
        sys.stdout = old_stdout
        help_text = captured.getvalue()

        arg_names = list(parser.arguments.keys())

        missing = []
        for arg_name in arg_names:
            if f"--{arg_name}" not in help_text:
                missing.append(arg_name)

        assert missing == [], f"Fields missing from help output: {missing}"

    def test_non_poly_enum_field_not_filtered(self):
        """Non-polymorphic enum fields (like attention_backend_type) must not
        be treated as type selectors and must appear in help output."""
        walker = ConfigWalker(InferenceEngineConfig)
        parser = walker.build_parser()

        # Find the attention_backend_type field
        abt_name = None
        for name in parser.arguments:
            if "attention_backend_type" in name:
                abt_name = name
                break

        assert abt_name is not None, "attention_backend_type not found in parser"
        arg = parser.arguments[abt_name]
        assert (
            not arg.is_type_selector
        ), "attention_backend_type should NOT be a type selector"
        assert (
            arg.poly_base_field is None
        ), "attention_backend_type should NOT have a poly_base_field"

    def test_non_poly_enum_field_in_correct_group(self):
        """Non-poly enum fields should be grouped with their parent config,
        not sent to the root group."""
        walker = ConfigWalker(InferenceEngineConfig)
        parser = walker.build_parser()

        abt_arg = None
        for name, arg in parser.arguments.items():
            if "attention_backend_type" in name:
                abt_arg = arg
                break

        assert abt_arg is not None
        # Should be in a nested group, not root ("")
        assert (
            abt_arg.group != ""
        ), f"attention_backend_type in root group, expected nested group"

    def test_type_selector_has_correct_attributes(self):
        """Actual polymorphic type selectors must have is_type_selector=True
        and a valid poly_base_field."""
        walker = ConfigWalker(InferenceEngineConfig)
        parser = walker.build_parser()

        # controller_config.type is a real type selector
        controller_type = parser.arguments.get("controller_config.type")
        assert controller_type is not None
        assert controller_type.is_type_selector is True
        assert controller_type.poly_base_field == "controller_config"
        assert controller_type.choices is not None
        assert "llm" in controller_type.choices

    def test_metrics_not_classified_as_polymorphic_group(self):
        """MetricsConfig is not polymorphic — the metrics_config node in the
        help tree should be a plain nested DC group, not a poly node."""
        walker = ConfigWalker(InferenceEngineConfig)
        parser = walker.build_parser()

        metrics_node = parser.help_tree.groups.get("metrics_config")
        assert metrics_node is not None, "metrics_config should be in help tree"
        assert (
            metrics_node.type_arg is None and not metrics_node.variants
        ), "metrics_config should NOT be a polymorphic node"

    def test_metrics_fields_in_help_output(self):
        """All MetricsConfig fields must appear in help output, including
        the deeply nested schedule.type selector."""
        walker = ConfigWalker(InferenceEngineConfig)
        parser = walker.build_parser()

        old_stdout = sys.stdout
        sys.stdout = captured = StringIO()
        parser.print_help()
        sys.stdout = old_stdout
        help_text = captured.getvalue()

        expected_metrics_args = [
            "metrics_config.write_metrics",
            "metrics_config.wandb.enabled",
            "metrics_config.wandb.project",
            "metrics_config.chrome_trace.enabled",
            "metrics_config.chrome_trace.enable_cpu_ops",
            "metrics_config.gpu_profiling.enabled",
            "metrics_config.gpu_profiling.output_dir",
            "metrics_config.gpu_profiling.schedule.type",
            "metrics_config.gpu_profiling.schedule.traces_per_batch_shape",
            "metrics_config.gpu_profiling.schedule.warmup_batches",
            "metrics_config.gpu_profiling.schedule.batches_per_cycle",
            "metrics_config.gpu_profiling.schedule.num_cycles",
            "metrics_config.enable_cpu_op_level",
            "metrics_config.enable_per_batch_metrics",
            "metrics_config.enable_png_output",
        ]

        missing = [arg for arg in expected_metrics_args if f"--{arg}" not in help_text]
        assert missing == [], f"Metrics fields missing from help: {missing}"

    def test_type_selectors_shown_in_poly_groups_not_duplicated(self):
        """Type selectors should be stored as type_arg on poly HelpNodes
        and have is_type_selector=True with a poly_base_field."""
        walker = ConfigWalker(InferenceEngineConfig)
        parser = walker.build_parser()

        # Type selectors should have is_type_selector=True and a poly_base_field
        ct = parser.arguments.get("controller_config.type")
        assert ct is not None
        assert ct.is_type_selector is True
        assert ct.poly_base_field == "controller_config"

        rat = parser.arguments.get("resource_allocator_config.type")
        assert rat is not None
        assert rat.is_type_selector is True
        assert rat.poly_base_field == "resource_allocator_config"

        # Poly groups should be identified via help tree
        controller_node = parser.help_tree.groups.get("controller_config")
        assert controller_node is not None
        assert controller_node.type_arg is not None
        assert len(controller_node.variants) > 0


class TestNonPolyEnumRoundtrip:
    """Test serialization/deserialization of non-polymorphic enum fields."""

    def test_attention_backend_type_default(self):
        """Test default enum value survives roundtrip."""
        config = WorkerConfig()
        assert config.attention_backend_type == AttentionBackendType.FLASH_ATTENTION

        as_dict = dataclass_to_dict(config)
        assert as_dict["attention_backend_type"] == "flash_attention"

        reconstructed = create_class_from_dict(WorkerConfig, as_dict)
        assert (
            reconstructed.attention_backend_type == AttentionBackendType.FLASH_ATTENTION
        )

    def test_attention_backend_type_non_default(self):
        """Test non-default enum value survives roundtrip."""
        config = WorkerConfig(attention_backend_type=AttentionBackendType.FLASHINFER)
        as_dict = dataclass_to_dict(config)
        assert as_dict["attention_backend_type"] == "flashinfer"

        reconstructed = create_class_from_dict(WorkerConfig, as_dict)
        assert reconstructed.attention_backend_type == AttentionBackendType.FLASHINFER

    def test_attention_backend_in_nested_config(self):
        """Test enum field survives roundtrip when nested in full hierarchy."""
        config_dict = {
            "controller_config": {
                "type": "llm",
                "replica_controller_config": {
                    "type": "llm_base",
                    "worker_config": {
                        "attention_backend_type": "flashinfer",
                    },
                },
            },
        }
        config = create_class_from_dict(InferenceEngineConfig, config_dict)
        worker = config.controller_config.replica_controller_config.worker_config
        assert worker.attention_backend_type == AttentionBackendType.FLASHINFER

        # Full roundtrip
        as_dict = dataclass_to_dict(config)
        reconstructed = create_class_from_dict(InferenceEngineConfig, as_dict)
        r_worker = (
            reconstructed.controller_config.replica_controller_config.worker_config
        )
        assert r_worker.attention_backend_type == AttentionBackendType.FLASHINFER


class TestDeeplyNestedPolyInNonPolyConfig:
    """Test configs where a non-polymorphic config contains deeply nested
    polymorphic fields (MetricsConfig -> GpuProfilingConfig -> schedule)."""

    def test_metrics_default_instantiation(self):
        """Default MetricsConfig uses batchwise schedule."""
        config = MetricsConfig()
        assert isinstance(config.gpu_profiling.schedule, BatchwiseProfilingSchedule)
        assert config.gpu_profiling.schedule.traces_per_batch_shape == 3

    def test_metrics_with_cyclical_schedule(self):
        """MetricsConfig can use cyclical schedule."""
        config = MetricsConfig(
            gpu_profiling=GpuProfilingConfig(
                enabled=True,
                schedule=CyclicalProfilingSchedule(
                    warmup_batches=10, batches_per_cycle=50
                ),
            )
        )
        assert isinstance(config.gpu_profiling.schedule, CyclicalProfilingSchedule)
        assert config.gpu_profiling.schedule.warmup_batches == 10
        assert config.gpu_profiling.schedule.batches_per_cycle == 50

    def test_metrics_dict_roundtrip(self):
        """MetricsConfig roundtrip through dict preserves schedule variant."""
        original = MetricsConfig(
            write_metrics=True,
            wandb=WandbConfig(enabled=True, project="test_proj"),
            gpu_profiling=GpuProfilingConfig(
                enabled=True,
                schedule=CyclicalProfilingSchedule(
                    warmup_batches=5, batches_per_cycle=20, num_cycles=3
                ),
            ),
            enable_cpu_op_level=True,
        )

        as_dict = dataclass_to_dict(original)
        assert as_dict["write_metrics"] is True
        assert as_dict["wandb"]["enabled"] is True
        assert as_dict["wandb"]["project"] == "test_proj"
        assert as_dict["gpu_profiling"]["schedule"]["type"] == "cyclical"
        assert as_dict["gpu_profiling"]["schedule"]["warmup_batches"] == 5

        reconstructed = create_class_from_dict(MetricsConfig, as_dict)
        assert reconstructed.write_metrics is True
        assert reconstructed.wandb.enabled is True
        assert reconstructed.wandb.project == "test_proj"
        assert isinstance(
            reconstructed.gpu_profiling.schedule, CyclicalProfilingSchedule
        )
        assert reconstructed.gpu_profiling.schedule.warmup_batches == 5
        assert reconstructed.gpu_profiling.schedule.batches_per_cycle == 20
        assert reconstructed.gpu_profiling.schedule.num_cycles == 3
        assert reconstructed.enable_cpu_op_level is True

    def test_metrics_in_full_hierarchy_roundtrip(self):
        """MetricsConfig roundtrip within the full InferenceEngineConfig."""
        config_dict = {
            "metrics_config": {
                "write_metrics": True,
                "wandb": {"enabled": True, "project": "my_proj"},
                "chrome_trace": {"enabled": True, "enable_cpu_ops": True},
                "gpu_profiling": {
                    "enabled": True,
                    "output_dir": "/custom/traces",
                    "schedule": {
                        "type": "cyclical",
                        "warmup_batches": 10,
                        "batches_per_cycle": 100,
                        "num_cycles": 5,
                    },
                },
                "enable_cpu_op_level": True,
                "enable_per_batch_metrics": True,
                "enable_png_output": True,
            }
        }

        config = create_class_from_dict(InferenceEngineConfig, config_dict)
        m = config.metrics_config

        assert m.write_metrics is True
        assert m.wandb.enabled is True
        assert m.wandb.project == "my_proj"
        assert m.chrome_trace.enabled is True
        assert m.chrome_trace.enable_cpu_ops is True
        assert m.gpu_profiling.enabled is True
        assert m.gpu_profiling.output_dir == "/custom/traces"
        assert isinstance(m.gpu_profiling.schedule, CyclicalProfilingSchedule)
        assert m.gpu_profiling.schedule.warmup_batches == 10
        assert m.gpu_profiling.schedule.batches_per_cycle == 100
        assert m.gpu_profiling.schedule.num_cycles == 5
        assert m.enable_cpu_op_level is True
        assert m.enable_per_batch_metrics is True
        assert m.enable_png_output is True

        # Full roundtrip
        as_dict = dataclass_to_dict(config)
        reconstructed = create_class_from_dict(InferenceEngineConfig, as_dict)
        rm = reconstructed.metrics_config

        assert rm.write_metrics is True
        assert rm.wandb.enabled is True
        assert rm.wandb.project == "my_proj"
        assert rm.chrome_trace.enabled is True
        assert rm.gpu_profiling.output_dir == "/custom/traces"
        assert isinstance(rm.gpu_profiling.schedule, CyclicalProfilingSchedule)
        assert rm.gpu_profiling.schedule.warmup_batches == 10
        assert rm.enable_per_batch_metrics is True
        assert rm.enable_png_output is True

    def test_schedule_type_field_in_parser(self):
        """The profiling schedule.type must exist as a proper type selector
        in the nested parser, with choices for batchwise and cyclical."""
        walker = ConfigWalker(InferenceEngineConfig)
        parser = walker.build_parser()

        schedule_type = None
        for name, arg in parser.arguments.items():
            if "schedule.type" in name and "metrics" in name:
                schedule_type = arg
                break

        assert schedule_type is not None, "schedule.type not found"
        assert schedule_type.is_type_selector is True
        assert "batchwise" in schedule_type.choices
        assert "cyclical" in schedule_type.choices


class TestFullHierarchyRoundtrip:
    """Roundtrip serialization test exercising every field in the config tree.

    This mirrors the real vajra InferenceEngineConfig with all 88+ fields
    to ensure nothing is lost during dict -> dataclass -> dict -> dataclass.
    """

    def test_all_fields_roundtrip_with_non_default_values(self):
        """Create a config with every field set to a non-default value,
        serialize to dict and back, and verify all values survive."""
        config_dict = {
            "controller_config": {
                "type": "llm",
                "replica_controller_config": {
                    "type": "llm_base",
                    "model_deployment_config": {
                        "type": "llm",
                        "model": "custom-model-v2",
                        "tensor_parallel_size": 8,
                        "pipeline_parallel_size": 2,
                        "max_model_len": 8192,
                        "seed": 42,
                        "trust_remote_code": False,
                        "load_format": "dummy",
                    },
                    "worker_config": {
                        "num_workers": 4,
                        "worker_use_ray": True,
                        "attention_backend_type": "flashinfer",
                    },
                    "cache_config": {
                        "page_size": 64,
                        "max_promotion_batch_size": 2000,
                        "max_demotion_batch_size": 1500,
                        "allocator": {
                            "type": "cost_aware_exp_decay",
                            "prefill_profiling_path": "/data/profiling.json",
                            "decay_rate": 0.7,
                            "min_cost": 1e-8,
                        },
                        "gpu_tier": {
                            "admission_watermark": 0.02,
                            "demotion_watermark": 0.15,
                        },
                        "cpu_tier": {
                            "admission_watermark": 0.03,
                            "demotion_watermark": 0.2,
                        },
                        "nvme_tier": {
                            "admission_watermark": 0.01,
                            "demotion_watermark": 0.1,
                        },
                        "enable_prefix_caching": False,
                    },
                    "compute_manager_config": {
                        "type": "dynamic_chunk",
                        "min_chunk_size": 256,
                        "max_chunk_size": 8192,
                    },
                },
                "session_prioritizer_config": {
                    "type": "edf",
                    "default_deadline_ms": 500.0,
                },
                "session_router_config": {
                    "type": "round_robin",
                },
            },
            "metrics_config": {
                "write_metrics": True,
                "wandb": {"enabled": True, "project": "roundtrip_test"},
                "chrome_trace": {"enabled": True, "enable_cpu_ops": True},
                "gpu_profiling": {
                    "enabled": True,
                    "output_dir": "/output/traces",
                    "schedule": {
                        "type": "cyclical",
                        "warmup_batches": 15,
                        "batches_per_cycle": 50,
                        "num_cycles": 10,
                    },
                },
                "enable_cpu_op_level": True,
                "enable_per_batch_metrics": True,
                "enable_png_output": True,
            },
            "resource_allocator_config": {
                "type": "mock",
                "num_mock_gpus": 16,
            },
        }

        # First construction
        config = create_class_from_dict(InferenceEngineConfig, config_dict)

        # Serialize to dict
        as_dict = dataclass_to_dict(config)

        # Reconstruct
        reconstructed = create_class_from_dict(InferenceEngineConfig, as_dict)

        # ── controller_config ──
        cc = reconstructed.controller_config
        assert isinstance(cc, ReplicasetControllerConfig)

        # ── model_deployment_config ──
        md = cc.replica_controller_config.model_deployment_config
        assert md.model == "custom-model-v2"
        assert md.tensor_parallel_size == 8
        assert md.pipeline_parallel_size == 2
        assert md.max_model_len == 8192
        assert md.seed == 42
        assert md.trust_remote_code is False
        assert md.load_format == "dummy"

        # ── worker_config (with non-poly enum) ──
        wc = cc.replica_controller_config.worker_config
        assert wc.num_workers == 4
        assert wc.worker_use_ray is True
        assert wc.attention_backend_type == AttentionBackendType.FLASHINFER

        # ── cache_config ──
        cache = cc.replica_controller_config.cache_config
        assert cache.page_size == 64
        assert cache.max_promotion_batch_size == 2000
        assert cache.max_demotion_batch_size == 1500
        assert cache.enable_prefix_caching is False

        # ── allocator (polymorphic) ──
        alloc = cache.allocator
        assert isinstance(alloc, CostAwareExpDecayAllocatorConfig)
        assert alloc.prefill_profiling_path == "/data/profiling.json"
        assert alloc.decay_rate == 0.7
        assert alloc.min_cost == 1e-8

        # ── cache tiers ──
        assert cache.gpu_tier.admission_watermark == 0.02
        assert cache.gpu_tier.demotion_watermark == 0.15
        assert cache.cpu_tier.admission_watermark == 0.03
        assert cache.cpu_tier.demotion_watermark == 0.2
        assert cache.nvme_tier.admission_watermark == 0.01
        assert cache.nvme_tier.demotion_watermark == 0.1

        # ── compute_manager (polymorphic) ──
        cm = cc.replica_controller_config.compute_manager_config
        assert isinstance(cm, DynamicChunkComputeManagerConfig)
        assert cm.min_chunk_size == 256
        assert cm.max_chunk_size == 8192

        # ── session_prioritizer (polymorphic) ──
        sp = cc.session_prioritizer_config
        assert isinstance(sp, EdfSessionPrioritizerConfig)
        assert sp.default_deadline_ms == 500.0

        # ── session_router (polymorphic) ──
        sr = cc.session_router_config
        assert isinstance(sr, RoundRobinSessionRouterConfig)

        # ── metrics_config (non-poly with deeply nested poly) ──
        m = reconstructed.metrics_config
        assert m.write_metrics is True
        assert m.wandb.enabled is True
        assert m.wandb.project == "roundtrip_test"
        assert m.chrome_trace.enabled is True
        assert m.chrome_trace.enable_cpu_ops is True
        assert m.gpu_profiling.enabled is True
        assert m.gpu_profiling.output_dir == "/output/traces"
        assert isinstance(m.gpu_profiling.schedule, CyclicalProfilingSchedule)
        assert m.gpu_profiling.schedule.warmup_batches == 15
        assert m.gpu_profiling.schedule.batches_per_cycle == 50
        assert m.gpu_profiling.schedule.num_cycles == 10
        assert m.enable_cpu_op_level is True
        assert m.enable_per_batch_metrics is True
        assert m.enable_png_output is True

        # ── resource_allocator (polymorphic) ──
        ra = reconstructed.resource_allocator_config
        assert isinstance(ra, MockResourceAllocatorConfig)
        assert ra.num_mock_gpus == 16

    def _parse_cli(self, config_class, args):
        """Helper to parse CLI args by temporarily replacing sys.argv."""
        original_argv = sys.argv
        sys.argv = ["test"] + args
        try:
            return parse_cli_args(config_class)
        finally:
            sys.argv = original_argv

    def test_cli_roundtrip_metrics_fields(self):
        """Parse CLI args for metrics (non-poly config with nested poly)."""
        config = self._parse_cli(
            MetricsConfig,
            [
                "--write_metrics",
                "true",
                "--wandb.enabled",
                "true",
                "--wandb.project",
                "cli_test",
                "--chrome_trace.enabled",
                "true",
                "--chrome_trace.enable_cpu_ops",
                "true",
                "--gpu_profiling.enabled",
                "true",
                "--gpu_profiling.output_dir",
                "/tmp/traces",
                "--gpu_profiling.schedule.type",
                "cyclical",
                "--gpu_profiling.schedule.warmup_batches",
                "7",
                "--gpu_profiling.schedule.batches_per_cycle",
                "25",
                "--gpu_profiling.schedule.num_cycles",
                "3",
                "--enable_cpu_op_level",
                "true",
                "--enable_per_batch_metrics",
                "true",
                "--enable_png_output",
                "true",
            ],
        )

        assert config.write_metrics is True
        assert config.wandb.enabled is True
        assert config.wandb.project == "cli_test"
        assert config.chrome_trace.enabled is True
        assert config.chrome_trace.enable_cpu_ops is True
        assert config.gpu_profiling.enabled is True
        assert config.gpu_profiling.output_dir == "/tmp/traces"
        assert isinstance(config.gpu_profiling.schedule, CyclicalProfilingSchedule)
        assert config.gpu_profiling.schedule.warmup_batches == 7
        assert config.gpu_profiling.schedule.batches_per_cycle == 25
        assert config.gpu_profiling.schedule.num_cycles == 3
        assert config.enable_cpu_op_level is True
        assert config.enable_per_batch_metrics is True
        assert config.enable_png_output is True

    def test_cli_roundtrip_worker_with_enum(self):
        """Parse CLI args for worker config with non-poly enum field."""
        config = self._parse_cli(
            WorkerConfig,
            [
                "--num_workers",
                "4",
                "--worker_use_ray",
                "true",
                "--attention_backend_type",
                "flashinfer",
            ],
        )

        assert config.num_workers == 4
        assert config.worker_use_ray is True
        assert config.attention_backend_type == AttentionBackendType.FLASHINFER

    def test_cli_roundtrip_cache_config_with_poly_allocator(self):
        """Parse CLI args for a config with nested polymorphic fields."""
        config = self._parse_cli(
            CacheConfig,
            [
                "--page_size",
                "64",
                "--enable_prefix_caching",
                "false",
                "--allocator.type",
                "cost_aware_gdsf",
                "--allocator.use_frequency",
                "false",
                "--gpu_tier.admission_watermark",
                "0.02",
                "--gpu_tier.demotion_watermark",
                "0.15",
            ],
        )

        assert config.page_size == 64
        assert config.enable_prefix_caching is False
        assert isinstance(config.allocator, CostAwareGdsfAllocatorConfig)
        assert config.allocator.use_frequency is False
        assert config.gpu_tier.admission_watermark == 0.02
        assert config.gpu_tier.demotion_watermark == 0.15


# =============================================================================
# Help Rendering Tests
# =============================================================================
class TestHelpRendering:
    """Functional tests for rendered help output.

    These tests capture the actual stdout from print_help() and verify
    defaults, docstrings, variant organization, and absence of internal
    class names. They guard against bugs that structural metadata tests miss:

    Bug 1: BasePolyConfig() leaking into output instead of real docstrings
    Bug 2: Wrong per-variant defaults (all variants showing default variant's values)
    Bug 3: "only for" annotations appearing inside variant sections
    """

    @staticmethod
    def _capture_help(config_cls) -> str:
        """Build a parser for config_cls and capture its help output."""
        walker = ConfigWalker(config_cls)
        parser = walker.build_parser()
        old_stdout = sys.stdout
        sys.stdout = captured = StringIO()
        parser.print_help()
        sys.stdout = old_stdout
        return captured.getvalue()

    # ------------------------------------------------------------------
    # 1. Non-poly nested DCs are plain groups (no type_arg, no variants)
    # ------------------------------------------------------------------
    def test_help_tree_structure_plain_dc(self):
        """Non-poly nested DCs appear as plain HelpNode groups."""
        walker = ConfigWalker(TrainingConfig)
        parser = walker.build_parser()
        tree = parser.help_tree

        for name in ("parallelism", "batch", "memory", "paths"):
            node = tree.groups.get(name)
            assert node is not None, f"{name} not in help tree groups"
            assert node.type_arg is None, f"{name} should not have type_arg"
            assert not node.variants, f"{name} should not have variants"

    # ------------------------------------------------------------------
    # 2. Poly groups have type_arg, variants, and default_variant
    # ------------------------------------------------------------------
    def test_help_tree_structure_poly(self):
        """Poly groups have type_arg, populated variants, and default_variant."""
        walker = ConfigWalker(TrainingConfig)
        parser = walker.build_parser()
        tree = parser.help_tree

        # Scheduler
        sched = tree.groups.get("scheduler")
        assert sched is not None, "scheduler not in help tree"
        assert sched.type_arg is not None, "scheduler missing type_arg"
        assert set(sched.variants.keys()) == {"greedy", "balanced", "priority"}
        assert sched.default_variant == "greedy"

        # Storage
        storage = tree.groups.get("storage")
        assert storage is not None, "storage not in help tree"
        assert storage.type_arg is not None, "storage missing type_arg"
        assert set(storage.variants.keys()) == {"local", "s3", "gcs"}
        assert storage.default_variant == "local"

    # ------------------------------------------------------------------
    # 3. Variant-specific defaults in rendered output (catches Bug 2)
    # ------------------------------------------------------------------
    def test_variant_specific_defaults_in_output(self):
        """Each variant section shows its own fields with correct defaults."""
        help_text = self._capture_help(TrainingConfig)

        # balanced variant fields
        assert "balance_interval_ms" in help_text
        assert "max_imbalance_ratio" in help_text

        # priority variant fields
        assert "num_priority_levels" in help_text
        assert "preemption_enabled" in help_text
        assert "starvation_timeout_ms" in help_text

        # Check that balanced defaults appear: 100 and 0.2
        # Find the balanced section and verify defaults there
        walker = ConfigWalker(TrainingConfig)
        parser = walker.build_parser()
        tree = parser.help_tree
        balanced_node = tree.groups["scheduler"].variants["balanced"]

        balanced_defaults = {a.cli_name: a.default for a in balanced_node.args}
        assert balanced_defaults.get("scheduler.balance_interval_ms") == 100
        assert balanced_defaults.get("scheduler.max_imbalance_ratio") == 0.2

        # priority variant defaults
        priority_node = tree.groups["scheduler"].variants["priority"]
        priority_defaults = {a.cli_name: a.default for a in priority_node.args}
        assert priority_defaults.get("scheduler.num_priority_levels") == 3
        assert priority_defaults.get("scheduler.preemption_enabled") is True
        assert priority_defaults.get("scheduler.starvation_timeout_ms") == 5000

    # ------------------------------------------------------------------
    # 4. Shared field defaults correct per variant (catches Bug 2)
    # ------------------------------------------------------------------
    def test_shared_field_defaults_per_variant(self):
        """All scheduler variants inherit max_pending_requests with default 1000."""
        walker = ConfigWalker(TrainingConfig)
        parser = walker.build_parser()
        tree = parser.help_tree

        for variant_name in ("greedy", "balanced", "priority"):
            variant_node = tree.groups["scheduler"].variants[variant_name]
            defaults = {a.cli_name: a.default for a in variant_node.args}
            assert defaults.get("scheduler.max_pending_requests") == 1000, (
                f"{variant_name} variant has wrong max_pending_requests default: "
                f"{defaults.get('scheduler.max_pending_requests')}"
            )

    # ------------------------------------------------------------------
    # 5. No internal class names in help (catches Bug 1)
    # ------------------------------------------------------------------
    def test_no_internal_class_names_in_help(self):
        """Internal class names must not appear as raw text in help output."""
        training_help = self._capture_help(TrainingConfig)
        engine_help = self._capture_help(InferenceEngineConfig)

        forbidden = [
            "BasePolyConfig",
            "AbstractSchedulerConfig",
            "AbstractStorageBackendConfig",
            "AbstractReplicasetControllerConfig",
            "AbstractReplicaControllerConfig",
            "AbstractAllocatorConfig",
            "AbstractComputeManagerConfig",
            "AbstractSessionRouterConfig",
            "AbstractSessionPrioritizerConfig",
            "AbstractResourceAllocatorConfig",
            "AbstractModelDeploymentConfig",
            "AbstractProfilingSchedule",
        ]

        for name in forbidden:
            assert (
                name not in training_help
            ), f"Internal class name '{name}' found in TrainingConfig help"
            assert (
                name not in engine_help
            ), f"Internal class name '{name}' found in InferenceEngineConfig help"

    # ------------------------------------------------------------------
    # 6. Real docstrings appear in help (catches Bug 1)
    # ------------------------------------------------------------------
    def test_poly_docstrings_in_help(self):
        """Docstrings from abstract base and variant classes appear in help."""
        help_text = self._capture_help(TrainingConfig)

        # Abstract base docstring
        assert "Abstract base for scheduler configurations." in help_text

        # Variant docstrings
        assert "Greedy scheduler that processes requests immediately." in help_text
        assert "Balanced scheduler with load balancing." in help_text
        assert "Priority-based scheduler with multiple priority levels." in help_text

        # Storage docstrings
        assert "Abstract base for storage backend configurations." in help_text
        assert "Local filesystem storage." in help_text
        assert "AWS S3 storage backend." in help_text
        assert "Google Cloud Storage backend." in help_text

    # ------------------------------------------------------------------
    # 7. Type selector line with choices and default
    # ------------------------------------------------------------------
    def test_type_selector_line_in_help(self):
        """'Select variant with' line appears with correct choices and default."""
        help_text = self._capture_help(TrainingConfig)

        # Scheduler type selector
        assert "--scheduler.type" in help_text
        assert "balanced" in help_text
        assert "greedy" in help_text
        assert "priority" in help_text
        assert "Select variant with" in help_text

        # Storage type selector
        assert "--storage.type" in help_text

    # ------------------------------------------------------------------
    # 8. Default variant gets "(default)" marker
    # ------------------------------------------------------------------
    def test_default_variant_marker(self):
        """Default variant gets a '(default)' marker in help output."""
        help_text = self._capture_help(TrainingConfig)

        # greedy is the default scheduler variant
        # The output format is: "Scheduler → greedy (default)"
        # Look for "greedy" followed by "(default)" on the same line
        lines = help_text.split("\n")
        found_greedy_default = False
        for line in lines:
            if "greedy" in line and "(default)" in line:
                found_greedy_default = True
                break
        assert found_greedy_default, "greedy variant should have (default) marker"

        # local is the default storage variant
        found_local_default = False
        for line in lines:
            if "local" in line and "(default)" in line:
                found_local_default = True
                break
        assert found_local_default, "local variant should have (default) marker"

    # ------------------------------------------------------------------
    # 9. Variant fields only in correct variant section
    # ------------------------------------------------------------------
    def test_variant_fields_only_in_correct_variant(self):
        """Variant-specific fields only appear under their own variant."""
        walker = ConfigWalker(TrainingConfig)
        parser = walker.build_parser()
        tree = parser.help_tree

        # balance_interval_ms should be in balanced only
        for vname, vnode in tree.groups["scheduler"].variants.items():
            field_names = {a.cli_name for a in vnode.args}
            if vname == "balanced":
                assert "scheduler.balance_interval_ms" in field_names
                assert "scheduler.max_imbalance_ratio" in field_names
            else:
                assert "scheduler.balance_interval_ms" not in field_names
                assert "scheduler.max_imbalance_ratio" not in field_names

        # num_priority_levels should be in priority only
        for vname, vnode in tree.groups["scheduler"].variants.items():
            field_names = {a.cli_name for a in vnode.args}
            if vname == "priority":
                assert "scheduler.num_priority_levels" in field_names
                assert "scheduler.preemption_enabled" in field_names
                assert "scheduler.starvation_timeout_ms" in field_names
            else:
                assert "scheduler.num_priority_levels" not in field_names

        # Storage: root_path in local only, bucket in s3 and gcs
        for vname, vnode in tree.groups["storage"].variants.items():
            field_names = {a.cli_name for a in vnode.args}
            if vname == "local":
                assert "storage.root_path" in field_names
            else:
                assert "storage.root_path" not in field_names
            if vname in ("s3", "gcs"):
                assert "storage.bucket" in field_names
            else:
                assert "storage.bucket" not in field_names

    # ------------------------------------------------------------------
    # 10. Nested DC group headers in help output
    # ------------------------------------------------------------------
    def test_nested_dc_headers_in_help(self):
        """Nested DC group headers render correctly."""
        help_text = self._capture_help(TrainingConfig)

        assert "Parallelism Options:" in help_text
        assert "Batch Options:" in help_text
        assert "Memory Options:" in help_text
        assert "Paths Options:" in help_text

    # ------------------------------------------------------------------
    # 11. Deeply nested poly defaults (catches Bug 2 in deep nesting)
    # ------------------------------------------------------------------
    def test_deeply_nested_poly_defaults(self):
        """Default model appears under the correct variant in InferenceEngineConfig."""
        walker = ConfigWalker(InferenceEngineConfig)
        parser = walker.build_parser()
        tree = parser.help_tree

        # controller_config is a poly group; default variant is "llm"
        controller = tree.groups.get("controller_config")
        assert controller is not None
        assert controller.default_variant == "llm"

        llm_variant = controller.variants.get("llm")
        assert llm_variant is not None

        # Inside llm variant, replica_controller_config is a poly group
        replica = llm_variant.groups.get("replica_controller_config")
        assert replica is not None
        assert replica.default_variant == "llm_base"

        llm_base = replica.variants.get("llm_base")
        assert llm_base is not None

        # Inside llm_base variant, model_deployment_config is a poly group
        model_dep = llm_base.groups.get("model_deployment_config")
        assert model_dep is not None
        assert model_dep.default_variant == "llm"

        llm_model = model_dep.variants.get("llm")
        assert llm_model is not None

        # Verify default model value
        model_defaults = {a.cli_name: a.default for a in llm_model.args}
        model_field = None
        for cli_name, default in model_defaults.items():
            if cli_name.endswith(".model"):
                model_field = (cli_name, default)
                break
        assert model_field is not None, "model field not found in llm variant"
        assert (
            model_field[1] == "meta-llama/Llama-3-8B"
        ), f"Expected 'meta-llama/Llama-3-8B', got '{model_field[1]}'"

    # ------------------------------------------------------------------
    # 12. All fields present for TrainingConfig
    # ------------------------------------------------------------------
    def test_all_fields_present_training_config(self):
        """Every field from TrainingConfig hierarchy appears in help output."""
        help_text = self._capture_help(TrainingConfig)

        walker = ConfigWalker(TrainingConfig)
        parser = walker.build_parser()
        arg_names = list(parser.arguments.keys())

        missing = [name for name in arg_names if f"--{name}" not in help_text]
        assert (
            missing == []
        ), f"Fields missing from TrainingConfig help output: {missing}"
